lint and other fixes

Author: Czechh

datadog_lambda/tracing.py

@@ -155,11 +155,7 @@ def extract_context_custom_extractor(extractor, event, lambda_context):
     Extract Datadog trace context using a custom trace extractor function
     """
     try:
-        (
-            trace_id,
-            parent_id,
-            sampling_priority,
-        ) = extractor(event, lambda_context)
+        (trace_id, parent_id, sampling_priority,) = extractor(event, lambda_context)
         return trace_id, parent_id, sampling_priority
     except Exception as e:
         logger.debug("The trace extractor returned with error %s", e)
@@ -181,11 +177,9 @@ def extract_dd_trace_context(event, lambda_context, extractor=None):
     global dd_trace_context
 
     if extractor is not None:
-        (
-            trace_id,
-            parent_id,
-            sampling_priority,
-        ) = extract_context_custom_extractor(extractor, event, lambda_context)
+        (trace_id, parent_id, sampling_priority,) = extract_context_custom_extractor(
+            extractor, event, lambda_context
+        )
     elif "headers" in event:
         (
             trace_id,

tests/test_wrapper.py

@@ -105,28 +105,6 @@ def lambda_handler(event, context):
         self.mock_inject_correlation_ids.assert_called()
         self.mock_patch_all.assert_called()
 
-    def test_datadog_lambda_wrapper_with_extractor(self):
-        @datadog_lambda_wrapper
-        def lambda_handler(event, context):
-            lambda_metric("test.metric", 100)
-
-        lambda_event = {}
-
-        lambda_context = get_mock_context()
-
-        lambda_handler(lambda_event, lambda_context)
-
-        self.mock_metric_lambda_stats.distribution.assert_has_calls(
-            [call("test.metric", 100, timestamp=None, tags=ANY)]
-        )
-        self.mock_wrapper_lambda_stats.flush.assert_called()
-        self.mock_extract_dd_trace_context.assert_called_with(
-            lambda_event, lambda_context, extractor=None
-        )
-        self.mock_set_correlation_ids.assert_called()
-        self.mock_inject_correlation_ids.assert_called()
-        self.mock_patch_all.assert_called()
-
     def test_datadog_lambda_wrapper_flush_to_log(self):
         os.environ["DD_FLUSH_TO_LOG"] = "True"