Fix tests

DarcyRaynerDD · DarcyRaynerDD · commit 082683c5bbfb · 2020-03-24T13:55:09.000-04:00
diff --git a/datadog_lambda/tracing.py b/datadog_lambda/tracing.py
@@ -71,9 +71,9 @@ def _get_dd_trace_native_context():
     parent_id = span.context.span_id
     trace_id = span.context.trace_id
     return {
-        "parent_id": str(parent_id),
-        "trace_id": str(trace_id),
-        "sampling_priority": SamplingPriority.AUTO_KEEP,
+        "parent-id": str(parent_id),
+        "trace-id": str(trace_id),
+        "sampling-priority": SamplingPriority.AUTO_KEEP,
         "source": Source.DDTRACE,
     }
 
@@ -106,17 +106,17 @@ def extract_dd_trace_context(event):
     sampling_priority = lowercase_headers.get(TraceHeader.SAMPLING_PRIORITY)
     if trace_id and parent_id and sampling_priority:
         logger.debug("Extracted Datadog trace context from headers")
-        dd_trace_context = {
+        metadata = {
             "trace-id": trace_id,
             "parent-id": parent_id,
             "sampling-priority": sampling_priority,
-            "source": Source.EVENT,
         }
         xray_recorder.begin_subsegment(XraySubsegment.NAME)
         subsegment = xray_recorder.current_subsegment()
-        subsegment.put_metadata(
-            XraySubsegment.KEY, dd_trace_context, XraySubsegment.NAMESPACE
-        )
+
+        subsegment.put_metadata(XraySubsegment.KEY, metadata, XraySubsegment.NAMESPACE)
+        dd_trace_context = metadata.copy()
+        dd_trace_context["source"] = Source.EVENT
         xray_recorder.end_subsegment()
     else:
         # AWS Lambda runtime caches global variables between invocations,
diff --git a/datadog_lambda/wrapper.py b/datadog_lambda/wrapper.py
@@ -22,7 +22,7 @@
     set_correlation_ids,
 )
 from datadog_lambda.constants import Source
-from ddtrace import tracer, Span
+from ddtrace import tracer
 from ddtrace.propagation.http import HTTPPropagator
 
 
diff --git a/setup.py b/setup.py
@@ -30,6 +30,7 @@
     install_requires=[
         "aws-xray-sdk==2.4.3",
         "datadog==0.32.0",
+        "ddtrace==0.31.0",
         "wrapt==1.11.2",
         "setuptools==42.0.2",
     ],
diff --git a/tests/Dockerfile b/tests/Dockerfile
@@ -8,4 +8,5 @@ WORKDIR /test
 
 # Install datadog-lambda with dev dependencies from local
 COPY . .
-RUN pip install .[dev]
+RUN pip install .[dev]
+RUN pip install --find-links=https://s3.amazonaws.com/pypi.datadoghq.com/trace-dev/index.html ddtrace==0.35.1.dev5+g1e11b2dd -t .[dev]

Original file line number	Diff line number	Diff line change
`@@ -22,7 +22,7 @@`
`22`	`22`	`set_correlation_ids,`
`23`	`23`	`)`
`24`	`24`	`from datadog_lambda.constants import Source`
`25`		`-from ddtrace import tracer, Span`
	`25`	`+from ddtrace import tracer`
`26`	`26`	`from ddtrace.propagation.http import HTTPPropagator`
`27`	`27`
`28`	`28`