set_correlation_ids should handle the empty context case (#334)

* set_correlation_ids handles empty trace context case

Author: joeyzhao2018

datadog_lambda/tracing.py

@@ -515,10 +515,12 @@ def get_dd_trace_context():
     automatically, but this function can be used to manually inject the trace
     context to an outgoing request.
     """
-    global dd_trace_context
+    if dd_tracing_enabled:
+        dd_trace_py_context = _get_dd_trace_py_context()
+        if dd_trace_py_context is not None:
+            return _context_obj_to_headers(dd_trace_py_context)
 
-    context = None
-    xray_context = None
+    global dd_trace_context
 
     try:
         xray_context = _get_xray_trace_context()  # xray (sub)segment
@@ -527,22 +529,17 @@ def get_dd_trace_context():
             "get_dd_trace_context couldn't read from segment from x-ray, with error %s"
             % e
         )
+    if not xray_context:
+        return {}
 
-    if xray_context and not dd_trace_context:
-        context = xray_context
-    elif xray_context and dd_trace_context:
-        context = dd_trace_context.copy()
-        context["parent-id"] = xray_context.get("parent-id")
-        logger.debug(
-            "Set parent id from xray trace context: %s", context.get("parent-id")
-        )
+    if not dd_trace_context:
+        return _context_obj_to_headers(xray_context)
 
-    if dd_tracing_enabled:
-        dd_trace_py_context = _get_dd_trace_py_context()
-        if dd_trace_py_context is not None:
-            context = dd_trace_py_context
+    context = dd_trace_context.copy()
+    context["parent-id"] = xray_context.get("parent-id")
+    logger.debug("Set parent id from xray trace context: %s", context.get("parent-id"))
 
-    return _context_obj_to_headers(context) if context is not None else {}
+    return _context_obj_to_headers(context)
 
 
 def set_correlation_ids():
@@ -561,6 +558,8 @@ def set_correlation_ids():
         return
 
     context = get_dd_trace_context()
+    if not context:
+        return
 
     span = tracer.trace("dummy.span")
     span.trace_id = int(context[TraceHeader.TRACE_ID])

tests/test_tracing.py

@@ -42,6 +42,32 @@
 
 event_samples = "tests/event_samples/"
 
+span_to_finish = None
+
+
+def _clean_up_span():
+    global span_to_finish
+    if span_to_finish is not None:
+        span_to_finish.finish()
+        span_to_finish = None
+
+
+def register_span(span):
+    global span_to_finish
+    _clean_up_span()
+    span_to_finish = span
+    return span
+
+
+def wrapped_span_creator(span_creator_func):
+    def result_func(*args, **kwargs):
+        return register_span(span_creator_func(*args, **kwargs))
+
+    return result_func
+
+
+create_inferred_span = wrapped_span_creator(create_inferred_span)
+
 
 class ClientContext(object):
     def __init__(self, custom=None):
@@ -482,6 +508,15 @@ def test_set_correlation_ids(self):
         span = tracer.current_span()
         self.assertEqual(span.trace_id, 123)
         self.assertEqual(span.span_id, 456)
+        span.finish()
+
+    def test_set_correlation_ids_handle_empty_trace_context(self):
+        # neither x-ray or ddtrace is used. no tracing context at all.
+        self.mock_get_dd_trace_context.return_value = {}
+        # no exception thrown
+        set_correlation_ids()
+        span = tracer.current_span()
+        self.assertIsNone(span)
 
 
 class TestFunctionSpanTags(unittest.TestCase):
@@ -587,6 +622,9 @@ def setUp(self):
         self.mock_span_stop = patcher.start()
         self.addCleanup(patcher.stop)
 
+    def tearDown(self):
+        _clean_up_span()
+
     def test_create_inferred_span_from_authorizer_request_api_gateway_v1_event(self):
         event_sample_source = "authorizer-request-api-gateway-v1"
         finish_time = (
@@ -733,6 +771,9 @@ def _basic_common_checks(
 
 
 class TestInferredSpans(unittest.TestCase):
+    def tearDown(self):
+        _clean_up_span()
+
     def test_create_inferred_span_from_api_gateway_event(self):
         event_sample_source = "api-gateway"
         test_file = event_samples + event_sample_source + ".json"