Fix patching for integration tests

Author: tianchu

datadog_lambda/patch.py

@@ -10,8 +10,12 @@
 
 from wrapt import wrap_function_wrapper as wrap
 from wrapt.importer import when_imported
+from ddtrace import patch_all as patch_all_dd
 
-from datadog_lambda.tracing import get_dd_trace_context
+from datadog_lambda.tracing import (
+    get_dd_trace_context,
+    dd_tracing_enabled,
+)
 
 logger = logging.getLogger(__name__)
 
@@ -24,15 +28,31 @@
 
 _httplib_patched = False
 _requests_patched = False
+_integration_tests_patched = False
 
 
 def patch_all():
     """
-    Patch the widely-used HTTP clients to automatically inject
-    Datadog trace context.
+    Patch third-party libraries for tracing.
     """
-    _patch_httplib()
-    _ensure_patch_requests()
+    _patch_for_integration_tests()
+
+    if dd_tracing_enabled:
+        patch_all_dd()
+    else:
+        _patch_httplib()
+        _ensure_patch_requests()
+
+
+def _patch_for_integration_tests():
+    """
+    Patch `requests` to log the outgoing requests for integration tests.
+    """
+    global _integration_tests_patched
+    is_in_tests = os.environ.get("DD_INTEGRATION_TEST", "false").lower() == "true"
+    if not _integration_tests_patched and is_in_tests:
+        wrap("requests", "Session.send", _log_request)
+        _integration_tests_patched = True
 
 
 def _patch_httplib():
@@ -89,10 +109,6 @@ def _wrap_requests_request(func, instance, args, kwargs):
     else:
         kwargs["headers"] = context
 
-    # If we're in an integration test, log the HTTP requests made
-    if os.environ.get("DD_INTEGRATION_TEST", "false").lower() == "true":
-        _print_request_string(args, kwargs)
-
     return func(*args, **kwargs)
 
 
@@ -112,33 +128,28 @@ def _wrap_httplib_request(func, instance, args, kwargs):
     return func(*args, **kwargs)
 
 
-def _print_request_string(args, kwargs):
+def _log_request(func, instance, args, kwargs):
+    request = kwargs.get("request") or args[0]
+    _print_request_string(request)
+    return func(*args, **kwargs)
+
+
+def _print_request_string(request):
     """Print the request so that it can be checked in integration tests
 
     Only used by integration tests.
     """
-    # Normalizes the different ways args can be passed to a request
-    # to prevent test flakiness
-    method = None
-    if len(args) > 0:
-        method = args[0]
-    else:
-        method = kwargs.get("method", "").upper()
-
-    url = None
-    if len(args) > 1:
-        url = args[1]
-    else:
-        url = kwargs.get("url")
+    method = request.method
+    url = request.url
 
     # Sort the datapoints POSTed by their name so that snapshots always align
-    data = kwargs.get("data", "{}")
+    data = request.body or "{}"
     data_dict = json.loads(data)
     data_dict.get("series", []).sort(key=lambda series: series.get("metric"))
     sorted_data = json.dumps(data_dict)
 
     # Sort headers to prevent any differences in ordering
-    headers = kwargs.get("headers", {})
+    headers = request.headers or {}
     sorted_headers = sorted(
         "{}:{}".format(key, value) for key, value in headers.items()
     )

datadog_lambda/wrapper.py

@@ -22,7 +22,6 @@
     set_dd_trace_py_root,
     create_function_execution_span,
 )
-from ddtrace import patch_all as patch_all_dd
 
 logger = logging.getLogger(__name__)
 
@@ -93,12 +92,9 @@ def __init__(self, func):
             if self.logs_injection:
                 inject_correlation_ids()
 
-            if not dd_tracing_enabled:
-                # When using dd_trace_py it will patch all the http clients for us,
-                # Patch HTTP clients to propagate Datadog trace context
-                patch_all()
-            else:
-                patch_all_dd()
+            # Patch third-party libraries for tracing
+            patch_all()
+
             logger.debug("datadog_lambda_wrapper initialized")
         except Exception:
             traceback.print_exc()

tests/integration/snapshots/logs/http-requests_python27.log

@@ -1,17 +1,26 @@
 START RequestId: XXXX Version: $LATEST
 {"e": XXXX, "m": "aws.lambda.enhanced.invocations", "t": ["region:us-east-1", "account_id:XXXX", "functionname:integration-tester-dev-http-requests_python27", "cold_start:true", "memorysize:1024", "runtime:python2.7", "dd_lambda_layer:datadog-python27_2.XX.0"], "v": 1}
+HTTP GET https://ip-ranges.datadoghq.com/ Headers: ["Accept-Encoding:gzip, deflate", "Accept:*/*", "Connection:keep-alive", "User-Agent:python-requests/2.23.0", "x-datadog-parent-id:XXXX", "x-datadog-sampling-priority:1", "x-datadog-trace-id:XXXX"] Data: {}
+HTTP GET https://ip-ranges.datadoghq.eu/ Headers: ["Accept-Encoding:gzip, deflate", "Accept:*/*", "Connection:keep-alive", "User-Agent:python-requests/2.23.0", "x-datadog-parent-id:XXXX", "x-datadog-sampling-priority:1", "x-datadog-trace-id:XXXX"] Data: {}
+HTTP POST https://api.datadoghq.com/api/v1/distribution_points?api_key=XXXX Headers: ["Accept-Encoding:gzip, deflate", "Accept:*/*", "Connection:keep-alive", "Content-Length:460", "Content-Type:application/json", "User-Agent:python-requests/2.23.0", "x-datadog-parent-id:XXXX", "x-datadog-sampling-priority:1", "x-datadog-trace-id:XXXX"] Data: {"series": [{"tags": ["team:serverless", "role:hello", "dd_lambda_layer:datadog-python27_2.XX.0"], "metric": "hello.dog", "interval": 10, "host": null, "points": [[XXXX, [1.0]]], "device": null, "type": "distribution"}, {"tags": ["test:integration", "role:hello", "dd_lambda_layer:datadog-python27_2.XX.0"], "metric": "tests.integration.count", "interval": 10, "host": null, "points": [[XXXX, [21.0]]], "device": null, "type": "distribution"}]}
 {"traces": [[{"resource": "integration-tester-dev-http-requests_python27", "name": "aws.lambda", "service": "aws.lambda", "start": XXXX, "trace_id": "XXXX", "metrics": {"_sampling_priority_v1": 1, "system.pid": XXXX, "_dd.agent_psr": 1.0}, "parent_id": "XXXX", "meta": {"function_arn": "arn:aws:lambda:us-east-1:601427279990:function:integration-tester-dev-http-requests_python27", "resource_names": "integration-tester-dev-http-requests_python27", "cold_start": "true", "_dd.origin": "lambda", "_dd.parent_source": "xray", "request_id": "XXXX"}, "error": 0, "duration": XXXX, "type": "serverless", "span_id": "XXXX"}, {"resource": "requests.request", "name": "requests.request", "service": "aws.lambda", "start": XXXX, "trace_id": "XXXX", "metrics": {"_dd.measured": 1}, "parent_id": "XXXX", "meta": {"http.status_code": "200", "http.url": "https://ip-ranges.datadoghq.com/", "_dd.origin": "lambda", "http.method": "GET"}, "error": 0, "duration": XXXX, "type": "http", "span_id": "XXXX"}, {"resource": "requests.request", "name": "requests.request", "service": "aws.lambda", "start": XXXX, "trace_id": "XXXX", "metrics": {"_dd.measured": 1}, "parent_id": "XXXX", "meta": {"http.status_code": "200", "http.url": "https://ip-ranges.datadoghq.eu/", "_dd.origin": "lambda", "http.method": "GET"}, "error": 0, "duration": XXXX, "type": "http", "span_id": "XXXX"}, {"resource": "requests.request", "name": "requests.request", "service": "aws.lambda", "start": XXXX, "trace_id": "XXXX", "metrics": {"_dd.measured": 1}, "parent_id": "XXXX", "meta": {"http.status_code": "202", "http.url": "https://api.datadoghq.com/api/v1/distribution_points", "_dd.origin": "lambda", "http.method": "POST"}, "error": 0, "duration": XXXX, "type": "http", "span_id": "XXXX"}]]}
 END RequestId: XXXX
 REPORT RequestId: XXXX	Duration: XXXX ms	Billed Duration: XXXX ms	Memory Size: 1024 MB	Max Memory Used: XXXX MB	Init Duration: XXXX ms	
 XRAY TraceId: XXXX	SegmentId: XXXX	Sampled: true	
 START RequestId: XXXX Version: $LATEST
 {"e": XXXX, "m": "aws.lambda.enhanced.invocations", "t": ["region:us-east-1", "account_id:XXXX", "functionname:integration-tester-dev-http-requests_python27", "cold_start:false", "memorysize:1024", "runtime:python2.7", "dd_lambda_layer:datadog-python27_2.XX.0"], "v": 1}
+HTTP GET https://ip-ranges.datadoghq.com/ Headers: ["Accept-Encoding:gzip, deflate", "Accept:*/*", "Connection:keep-alive", "User-Agent:python-requests/2.23.0", "x-datadog-parent-id:XXXX", "x-datadog-sampling-priority:1", "x-datadog-trace-id:XXXX"] Data: {}
+HTTP GET https://ip-ranges.datadoghq.eu/ Headers: ["Accept-Encoding:gzip, deflate", "Accept:*/*", "Connection:keep-alive", "User-Agent:python-requests/2.23.0", "x-datadog-parent-id:XXXX", "x-datadog-sampling-priority:1", "x-datadog-trace-id:XXXX"] Data: {}
+HTTP POST https://api.datadoghq.com/api/v1/distribution_points?api_key=XXXX Headers: ["Accept-Encoding:gzip, deflate", "Accept:*/*", "Connection:keep-alive", "Content-Length:460", "Content-Type:application/json", "User-Agent:python-requests/2.23.0", "x-datadog-parent-id:XXXX", "x-datadog-sampling-priority:1", "x-datadog-trace-id:XXXX"] Data: {"series": [{"tags": ["team:serverless", "role:hello", "dd_lambda_layer:datadog-python27_2.XX.0"], "metric": "hello.dog", "interval": 10, "host": null, "points": [[XXXX, [1.0]]], "device": null, "type": "distribution"}, {"tags": ["test:integration", "role:hello", "dd_lambda_layer:datadog-python27_2.XX.0"], "metric": "tests.integration.count", "interval": 10, "host": null, "points": [[XXXX, [21.0]]], "device": null, "type": "distribution"}]}
 {"traces": [[{"resource": "integration-tester-dev-http-requests_python27", "name": "aws.lambda", "service": "aws.lambda", "start": XXXX, "trace_id": "XXXX", "metrics": {"_sampling_priority_v1": 1, "system.pid": XXXX, "_dd.agent_psr": 1.0}, "parent_id": "XXXX", "meta": {"function_arn": "arn:aws:lambda:us-east-1:601427279990:function:integration-tester-dev-http-requests_python27", "resource_names": "integration-tester-dev-http-requests_python27", "cold_start": "false", "_dd.origin": "lambda", "_dd.parent_source": "xray", "request_id": "XXXX"}, "error": 0, "duration": XXXX, "type": "serverless", "span_id": "XXXX"}, {"resource": "requests.request", "name": "requests.request", "service": "aws.lambda", "start": XXXX, "trace_id": "XXXX", "metrics": {"_dd.measured": 1}, "parent_id": "XXXX", "meta": {"http.status_code": "200", "http.url": "https://ip-ranges.datadoghq.com/", "_dd.origin": "lambda", "http.method": "GET"}, "error": 0, "duration": XXXX, "type": "http", "span_id": "XXXX"}, {"resource": "requests.request", "name": "requests.request", "service": "aws.lambda", "start": XXXX, "trace_id": "XXXX", "metrics": {"_dd.measured": 1}, "parent_id": "XXXX", "meta": {"http.status_code": "200", "http.url": "https://ip-ranges.datadoghq.eu/", "_dd.origin": "lambda", "http.method": "GET"}, "error": 0, "duration": XXXX, "type": "http", "span_id": "XXXX"}, {"resource": "requests.request", "name": "requests.request", "service": "aws.lambda", "start": XXXX, "trace_id": "XXXX", "metrics": {"_dd.measured": 1}, "parent_id": "XXXX", "meta": {"http.status_code": "202", "http.url": "https://api.datadoghq.com/api/v1/distribution_points", "_dd.origin": "lambda", "http.method": "POST"}, "error": 0, "duration": XXXX, "type": "http", "span_id": "XXXX"}]]}
 END RequestId: XXXX
 REPORT RequestId: XXXX	Duration: XXXX ms	Billed Duration: XXXX ms	Memory Size: 1024 MB	Max Memory Used: XXXX MB	
 XRAY TraceId: XXXX	SegmentId: XXXX	Sampled: true	
 START RequestId: XXXX Version: $LATEST
 {"e": XXXX, "m": "aws.lambda.enhanced.invocations", "t": ["region:us-east-1", "account_id:XXXX", "functionname:integration-tester-dev-http-requests_python27", "cold_start:false", "memorysize:1024", "runtime:python2.7", "dd_lambda_layer:datadog-python27_2.XX.0"], "v": 1}
+HTTP GET https://ip-ranges.datadoghq.com/ Headers: ["Accept-Encoding:gzip, deflate", "Accept:*/*", "Connection:keep-alive", "User-Agent:python-requests/2.23.0", "x-datadog-parent-id:XXXX", "x-datadog-sampling-priority:1", "x-datadog-trace-id:XXXX"] Data: {}
+HTTP GET https://ip-ranges.datadoghq.eu/ Headers: ["Accept-Encoding:gzip, deflate", "Accept:*/*", "Connection:keep-alive", "User-Agent:python-requests/2.23.0", "x-datadog-parent-id:XXXX", "x-datadog-sampling-priority:1", "x-datadog-trace-id:XXXX"] Data: {}
+HTTP POST https://api.datadoghq.com/api/v1/distribution_points?api_key=XXXX Headers: ["Accept-Encoding:gzip, deflate", "Accept:*/*", "Connection:keep-alive", "Content-Length:460", "Content-Type:application/json", "User-Agent:python-requests/2.23.0", "x-datadog-parent-id:XXXX", "x-datadog-sampling-priority:1", "x-datadog-trace-id:XXXX"] Data: {"series": [{"tags": ["team:serverless", "role:hello", "dd_lambda_layer:datadog-python27_2.XX.0"], "metric": "hello.dog", "interval": 10, "host": null, "points": [[XXXX, [1.0]]], "device": null, "type": "distribution"}, {"tags": ["test:integration", "role:hello", "dd_lambda_layer:datadog-python27_2.XX.0"], "metric": "tests.integration.count", "interval": 10, "host": null, "points": [[XXXX, [21.0]]], "device": null, "type": "distribution"}]}
 {"traces": [[{"resource": "integration-tester-dev-http-requests_python27", "name": "aws.lambda", "service": "aws.lambda", "start": XXXX, "trace_id": "XXXX", "metrics": {"_sampling_priority_v1": 1, "system.pid": XXXX, "_dd.agent_psr": 1.0}, "parent_id": "XXXX", "meta": {"function_arn": "arn:aws:lambda:us-east-1:601427279990:function:integration-tester-dev-http-requests_python27", "resource_names": "integration-tester-dev-http-requests_python27", "cold_start": "false", "_dd.origin": "lambda", "_dd.parent_source": "xray", "request_id": "XXXX"}, "error": 0, "duration": XXXX, "type": "serverless", "span_id": "XXXX"}, {"resource": "requests.request", "name": "requests.request", "service": "aws.lambda", "start": XXXX, "trace_id": "XXXX", "metrics": {"_dd.measured": 1}, "parent_id": "XXXX", "meta": {"http.status_code": "200", "http.url": "https://ip-ranges.datadoghq.com/", "_dd.origin": "lambda", "http.method": "GET"}, "error": 0, "duration": XXXX, "type": "http", "span_id": "XXXX"}, {"resource": "requests.request", "name": "requests.request", "service": "aws.lambda", "start": XXXX, "trace_id": "XXXX", "metrics": {"_dd.measured": 1}, "parent_id": "XXXX", "meta": {"http.status_code": "200", "http.url": "https://ip-ranges.datadoghq.eu/", "_dd.origin": "lambda", "http.method": "GET"}, "error": 0, "duration": XXXX, "type": "http", "span_id": "XXXX"}, {"resource": "requests.request", "name": "requests.request", "service": "aws.lambda", "start": XXXX, "trace_id": "XXXX", "metrics": {"_dd.measured": 1}, "parent_id": "XXXX", "meta": {"http.status_code": "202", "http.url": "https://api.datadoghq.com/api/v1/distribution_points", "_dd.origin": "lambda", "http.method": "POST"}, "error": 0, "duration": XXXX, "type": "http", "span_id": "XXXX"}]]}
 END RequestId: XXXX
 REPORT RequestId: XXXX	Duration: XXXX ms	Billed Duration: XXXX ms	Memory Size: 1024 MB	Max Memory Used: XXXX MB