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.request", _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,6 +128,11 @@ def _wrap_httplib_request(func, instance, args, kwargs):
     return func(*args, **kwargs)
 
 
+def _log_request(func, instance, args, kwargs):
+    _print_request_string(args, kwargs)
+    return func(*args, **kwargs)
+
+
 def _print_request_string(args, kwargs):
     """Print the request so that it can be checked in integration tests
 

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: [] Data: {}
+HTTP GET https://ip-ranges.datadoghq.eu/ Headers: [] Data: {}
+HTTP POST https://api.datadoghq.com/api/v1/distribution_points Headers: ["Content-Type:application/json"] 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: [] Data: {}
+HTTP GET https://ip-ranges.datadoghq.eu/ Headers: [] Data: {}
+HTTP POST https://api.datadoghq.com/api/v1/distribution_points Headers: ["Content-Type:application/json"] 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: [] Data: {}
+HTTP GET https://ip-ranges.datadoghq.eu/ Headers: [] Data: {}
+HTTP POST https://api.datadoghq.com/api/v1/distribution_points Headers: ["Content-Type:application/json"] 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	

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

@@ -1,17 +1,26 @@
 START RequestId: XXXX Version: $LATEST
 {"m": "aws.lambda.enhanced.invocations", "v": 1, "e": XXXX, "t": ["region:us-east-1", "account_id:XXXX", "functionname:integration-tester-dev-http-requests_python36", "cold_start:true", "memorysize:1024", "runtime:python3.6", "dd_lambda_layer:datadog-python36_2.XX.0"]}
+HTTP GET https://ip-ranges.datadoghq.com/ Headers: [] Data: {}
+HTTP GET https://ip-ranges.datadoghq.eu/ Headers: [] Data: {}
+HTTP POST https://api.datadoghq.com/api/v1/distribution_points Headers: ["Content-Type:application/json"] Data: {"series": [{"metric": "hello.dog", "points": [[XXXX, [1.0]]], "type": "distribution", "host": null, "device": null, "tags": ["team:serverless", "role:hello", "dd_lambda_layer:datadog-python36_2.XX.0"], "interval": 10}, {"metric": "tests.integration.count", "points": [[XXXX, [21.0]]], "type": "distribution", "host": null, "device": null, "tags": ["test:integration", "role:hello", "dd_lambda_layer:datadog-python36_2.XX.0"], "interval": 10}]}
 {"traces": [[{"trace_id": "XXXX", "parent_id": "XXXX", "span_id": "XXXX", "service": "aws.lambda", "resource": "integration-tester-dev-http-requests_python36", "name": "aws.lambda", "error": 0, "start": XXXX, "duration": XXXX, "meta": {"_dd.origin": "lambda", "cold_start": "true", "function_arn": "arn:aws:lambda:us-east-1:601427279990:function:integration-tester-dev-http-requests_python36", "request_id": "XXXX", "resource_names": "integration-tester-dev-http-requests_python36", "_dd.parent_source": "xray"}, "metrics": {"_dd.agent_psr": 1.0, "system.pid": XXXX, "_sampling_priority_v1": 1}, "type": "serverless"}, {"trace_id": "XXXX", "parent_id": "XXXX", "span_id": "XXXX", "service": "aws.lambda", "resource": "requests.request", "name": "requests.request", "error": 0, "start": XXXX, "duration": XXXX, "meta": {"_dd.origin": "lambda", "http.method": "GET", "http.url": "https://ip-ranges.datadoghq.com/", "http.status_code": "200"}, "metrics": {"_dd.measured": 1}, "type": "http"}, {"trace_id": "XXXX", "parent_id": "XXXX", "span_id": "XXXX", "service": "aws.lambda", "resource": "requests.request", "name": "requests.request", "error": 0, "start": XXXX, "duration": XXXX, "meta": {"_dd.origin": "lambda", "http.method": "GET", "http.url": "https://ip-ranges.datadoghq.eu/", "http.status_code": "200"}, "metrics": {"_dd.measured": 1}, "type": "http"}, {"trace_id": "XXXX", "parent_id": "XXXX", "span_id": "XXXX", "service": "aws.lambda", "resource": "requests.request", "name": "requests.request", "error": 0, "start": XXXX, "duration": XXXX, "meta": {"_dd.origin": "lambda", "http.method": "POST", "http.url": "https://api.datadoghq.com/api/v1/distribution_points", "http.status_code": "202"}, "metrics": {"_dd.measured": 1}, "type": "http"}]]}
 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
 {"m": "aws.lambda.enhanced.invocations", "v": 1, "e": XXXX, "t": ["region:us-east-1", "account_id:XXXX", "functionname:integration-tester-dev-http-requests_python36", "cold_start:false", "memorysize:1024", "runtime:python3.6", "dd_lambda_layer:datadog-python36_2.XX.0"]}
+HTTP GET https://ip-ranges.datadoghq.com/ Headers: [] Data: {}
+HTTP GET https://ip-ranges.datadoghq.eu/ Headers: [] Data: {}
+HTTP POST https://api.datadoghq.com/api/v1/distribution_points Headers: ["Content-Type:application/json"] Data: {"series": [{"metric": "hello.dog", "points": [[XXXX, [1.0]]], "type": "distribution", "host": null, "device": null, "tags": ["team:serverless", "role:hello", "dd_lambda_layer:datadog-python36_2.XX.0"], "interval": 10}, {"metric": "tests.integration.count", "points": [[XXXX, [21.0]]], "type": "distribution", "host": null, "device": null, "tags": ["test:integration", "role:hello", "dd_lambda_layer:datadog-python36_2.XX.0"], "interval": 10}]}
 {"traces": [[{"trace_id": "XXXX", "parent_id": "XXXX", "span_id": "XXXX", "service": "aws.lambda", "resource": "integration-tester-dev-http-requests_python36", "name": "aws.lambda", "error": 0, "start": XXXX, "duration": XXXX, "meta": {"_dd.origin": "lambda", "cold_start": "false", "function_arn": "arn:aws:lambda:us-east-1:601427279990:function:integration-tester-dev-http-requests_python36", "request_id": "XXXX", "resource_names": "integration-tester-dev-http-requests_python36", "_dd.parent_source": "xray"}, "metrics": {"_dd.agent_psr": 1.0, "system.pid": XXXX, "_sampling_priority_v1": 1}, "type": "serverless"}, {"trace_id": "XXXX", "parent_id": "XXXX", "span_id": "XXXX", "service": "aws.lambda", "resource": "requests.request", "name": "requests.request", "error": 0, "start": XXXX, "duration": XXXX, "meta": {"_dd.origin": "lambda", "http.method": "GET", "http.url": "https://ip-ranges.datadoghq.com/", "http.status_code": "200"}, "metrics": {"_dd.measured": 1}, "type": "http"}, {"trace_id": "XXXX", "parent_id": "XXXX", "span_id": "XXXX", "service": "aws.lambda", "resource": "requests.request", "name": "requests.request", "error": 0, "start": XXXX, "duration": XXXX, "meta": {"_dd.origin": "lambda", "http.method": "GET", "http.url": "https://ip-ranges.datadoghq.eu/", "http.status_code": "200"}, "metrics": {"_dd.measured": 1}, "type": "http"}, {"trace_id": "XXXX", "parent_id": "XXXX", "span_id": "XXXX", "service": "aws.lambda", "resource": "requests.request", "name": "requests.request", "error": 0, "start": XXXX, "duration": XXXX, "meta": {"_dd.origin": "lambda", "http.method": "POST", "http.url": "https://api.datadoghq.com/api/v1/distribution_points", "http.status_code": "202"}, "metrics": {"_dd.measured": 1}, "type": "http"}]]}
 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
 {"m": "aws.lambda.enhanced.invocations", "v": 1, "e": XXXX, "t": ["region:us-east-1", "account_id:XXXX", "functionname:integration-tester-dev-http-requests_python36", "cold_start:false", "memorysize:1024", "runtime:python3.6", "dd_lambda_layer:datadog-python36_2.XX.0"]}
+HTTP GET https://ip-ranges.datadoghq.com/ Headers: [] Data: {}
+HTTP GET https://ip-ranges.datadoghq.eu/ Headers: [] Data: {}
+HTTP POST https://api.datadoghq.com/api/v1/distribution_points Headers: ["Content-Type:application/json"] Data: {"series": [{"metric": "hello.dog", "points": [[XXXX, [1.0]]], "type": "distribution", "host": null, "device": null, "tags": ["team:serverless", "role:hello", "dd_lambda_layer:datadog-python36_2.XX.0"], "interval": 10}, {"metric": "tests.integration.count", "points": [[XXXX, [21.0]]], "type": "distribution", "host": null, "device": null, "tags": ["test:integration", "role:hello", "dd_lambda_layer:datadog-python36_2.XX.0"], "interval": 10}]}
 {"traces": [[{"trace_id": "XXXX", "parent_id": "XXXX", "span_id": "XXXX", "service": "aws.lambda", "resource": "integration-tester-dev-http-requests_python36", "name": "aws.lambda", "error": 0, "start": XXXX, "duration": XXXX, "meta": {"_dd.origin": "lambda", "cold_start": "false", "function_arn": "arn:aws:lambda:us-east-1:601427279990:function:integration-tester-dev-http-requests_python36", "request_id": "XXXX", "resource_names": "integration-tester-dev-http-requests_python36", "_dd.parent_source": "xray"}, "metrics": {"_dd.agent_psr": 1.0, "system.pid": XXXX, "_sampling_priority_v1": 1}, "type": "serverless"}, {"trace_id": "XXXX", "parent_id": "XXXX", "span_id": "XXXX", "service": "aws.lambda", "resource": "requests.request", "name": "requests.request", "error": 0, "start": XXXX, "duration": XXXX, "meta": {"_dd.origin": "lambda", "http.method": "GET", "http.url": "https://ip-ranges.datadoghq.com/", "http.status_code": "200"}, "metrics": {"_dd.measured": 1}, "type": "http"}, {"trace_id": "XXXX", "parent_id": "XXXX", "span_id": "XXXX", "service": "aws.lambda", "resource": "requests.request", "name": "requests.request", "error": 0, "start": XXXX, "duration": XXXX, "meta": {"_dd.origin": "lambda", "http.method": "GET", "http.url": "https://ip-ranges.datadoghq.eu/", "http.status_code": "200"}, "metrics": {"_dd.measured": 1}, "type": "http"}, {"trace_id": "XXXX", "parent_id": "XXXX", "span_id": "XXXX", "service": "aws.lambda", "resource": "requests.request", "name": "requests.request", "error": 0, "start": XXXX, "duration": XXXX, "meta": {"_dd.origin": "lambda", "http.method": "POST", "http.url": "https://api.datadoghq.com/api/v1/distribution_points", "http.status_code": "202"}, "metrics": {"_dd.measured": 1}, "type": "http"}]]}
 END RequestId: XXXX
 REPORT RequestId: XXXX	Duration: XXXX ms	Billed Duration: XXXX ms	Memory Size: 1024 MB	Max Memory Used: XXXX MB