test for correlation id in deferred tasks

update correlation_id if cid module installed

Fixes #119

Author: dshafer

django_declarative_apis/machinery/tasks.py

@@ -26,8 +26,10 @@
     import cid.locals
 
     _get_correlation_id = cid.locals.get_cid
+    _set_correlation_id = cid.locals.set_cid
 except ImportError:
     _get_correlation_id = lambda: None  # noqa: E731
+    _set_correlation_id = lambda _: None  # noqa: E731
 
 JOB_COUNT_CACHE_KEY = "future_task_runner:job_id"
 QUEUE_LENGTH_CACHE_KEY = "future_task_runner:current_queue_length"
@@ -151,6 +153,8 @@ def future_task_runner(
     resource_instance = resource_class.objects.get(pk=resource_instance_id)
     endpoint_task = getattr(endpoint_class, endpoint_method_name)
 
+    _set_correlation_id(correlation_id)
+
     _log_task_stats(
         endpoint_method_name,
         resource_instance_id,
@@ -291,6 +295,8 @@ def resource_task_runner(
     resource_instance = resource_class.objects.get(pk=resource_instance_id)
     resource_method = getattr(resource_instance, resource_method_name)
 
+    _set_correlation_id(correlation_id)
+
     _log_task_stats(
         "{0}.{1}".format(resource_class_name, resource_method_name),
         resource_instance_id,

tests/machinery/test_base.py

@@ -820,6 +820,38 @@ def setUp(self):
             "filtered_retry_count_2": 0,
         }
 
+    def test_future_task_runner_sets_cid(self):
+        data = {"cid": None}
+
+        def set_correlation_id(cid):
+            data["cid"] = cid
+
+        conf = tasks.future_task_runner.app.conf
+        old_val = conf["task_always_eager"]
+        conf["task_always_eager"] = True
+
+        old_set_cid = tasks._set_correlation_id
+        old_get_cid = tasks._get_correlation_id
+        tasks._set_correlation_id = set_correlation_id
+        tasks._get_correlation_id = lambda: "cid-sentinel"
+        try:
+            expected_response = {"foo": "bar"}
+            endpoint = _TestEndpoint(expected_response)
+            manager = machinery.EndpointBinder.BoundEndpointManager(
+                machinery._EndpointRequestLifecycleManager(endpoint), endpoint
+            )
+            machinery.EndpointBinder(endpoint).create_bound_endpoint(
+                manager, HttpRequest()
+            )
+
+            manager.get_response()
+        finally:
+            tasks._set_correlation_id = old_set_cid
+            tasks._get_correlation_id = old_get_cid
+            conf["task_always_eager"] = old_val
+
+        self.assertEqual("cid-sentinel", data["cid"])
+
     def test_get_response_kombu_error_retried(self):
         expected_response = {"foo": "bar"}
         endpoint = _TestEndpoint(expected_response)