diff --git a/datadog_lambda/wrapper.py b/datadog_lambda/wrapper.py index 8704eeab..e6d5d3b8 100644 --- a/datadog_lambda/wrapper.py +++ b/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) diff --git a/tests/test_wrapper.py b/tests/test_wrapper.py index 50f5e563..f383b63f 100644 --- a/tests/test_wrapper.py +++ b/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()