Fix wrapping for multiple handler functions in one file

Author: tianchu

datadog_lambda/wrapper.py

@@ -53,19 +53,19 @@ class _LambdaDecorator(object):
     and extracts/injects trace context.
     """
 
-    _instance = None
     _force_new = False
 
     def __new__(cls, func):
         """
-        If the decorator is accidentally applied to multiple functions or
-        the same function multiple times, only the first one takes effect.
+        If the decorator is accidentally applied to the same function multiple times,
+        only the first one takes effect.
 
         If _force_new, always return a real decorator, useful for unit tests.
         """
-        if cls._force_new or cls._instance is None:
-            cls._instance = super(_LambdaDecorator, cls).__new__(cls)
-            return cls._instance
+        if cls._force_new or not getattr(func, "_dd_wrapped", False):
+            wrapped = super(_LambdaDecorator, cls).__new__(cls)
+            wrapped._dd_wrapped = True
+            return wrapped
         else:
             return _NoopDecorator(func)
 

tests/test_wrapper.py

@@ -269,13 +269,11 @@ def lambda_handler(event, context):
         # the second @datadog_lambda_wrapper should actually be no-op.
         datadog_lambda_wrapper._force_new = False
 
-        @datadog_lambda_wrapper
-        def lambda_handler_wrapper(event, context):
-            lambda_handler(event, context)
+        lambda_handler_double_wrapped = datadog_lambda_wrapper(lambda_handler)
 
         lambda_event = {}
 
-        lambda_handler_wrapper(lambda_event, get_mock_context())
+        lambda_handler_double_wrapped(lambda_event, get_mock_context())
 
         self.mock_patch_all.assert_called_once()
         self.mock_wrapper_lambda_stats.flush.assert_called_once()