diff --git a/epsagon/wrappers/aws_lambda.py b/epsagon/wrappers/aws_lambda.py index b374da5c..014acd51 100644 --- a/epsagon/wrappers/aws_lambda.py +++ b/epsagon/wrappers/aws_lambda.py @@ -3,6 +3,7 @@ """ from __future__ import absolute_import +import os import traceback import time import copy @@ -35,6 +36,7 @@ def _add_status_code(runner, return_value): runner.resource['metadata']['status_code'] = status_code +# pylint: disable=too-many-statements def lambda_wrapper(func): """Epsagon's Lambda wrapper.""" @@ -58,6 +60,11 @@ def _lambda_wrapper(*args, **kwargs): # parameters / sends kwargs. In such case we ignore this trace. return func(*args, **kwargs) + if os.environ.get( + 'AWS_LAMBDA_INITIALIZATION_TYPE' + ) == 'provisioned-concurrency': + constants.COLD_START = False + try: runner = epsagon.runners.aws_lambda.LambdaRunner( time.time(), diff --git a/tests/wrappers/test_lambda_wrapper.py b/tests/wrappers/test_lambda_wrapper.py index 9e4df524..423af81c 100644 --- a/tests/wrappers/test_lambda_wrapper.py +++ b/tests/wrappers/test_lambda_wrapper.py @@ -1,3 +1,4 @@ +import os import json import mock import pytest @@ -81,9 +82,53 @@ def wrapped_lambda(event, context): trace_mock.add_exception.assert_not_called() assert not epsagon.constants.COLD_START + assert runner.resource['metadata']['cold_start'] == True assert runner.resource['metadata']['return_value'] == retval +@mock.patch.object(LambdaRunner, 'set_exception') +@mock.patch( + 'epsagon.trace.trace_factory.get_trace', + side_effect=lambda: trace_mock +) +@mock.patch( + 'epsagon.trace.trace_factory.get_or_create_trace', + side_effect=lambda: trace_mock +) +@mock.patch( + 'epsagon.triggers.aws_lambda.LambdaTriggerFactory.factory', + side_effect=['trigger'] +) +def test_lambda_wrapper_provisioned_concurrency_sanity( + trigger_factory_mock, + _, + __, + set_exception_mock +): + retval = 'success' + + @epsagon.wrappers.aws_lambda.lambda_wrapper + def wrapped_lambda(event, context): + return 'success' + + os.environ['AWS_LAMBDA_INITIALIZATION_TYPE'] = 'provisioned-concurrency' + assert wrapped_lambda('a', CONTEXT_STUB) == 'success' + trace_mock.prepare.assert_called() + runner = _get_runner_event(trace_mock) + + trigger_factory_mock.assert_called() + set_exception_mock.assert_not_called() + + trace_mock.set_timeout_handler.assert_called() + + trace_mock.send_traces.assert_called() + trace_mock.add_exception.assert_not_called() + + assert not epsagon.constants.COLD_START + assert runner.resource['metadata']['cold_start'] == False + os.environ.pop('AWS_LAMBDA_INITIALIZATION_TYPE') + + @mock.patch( 'epsagon.trace.trace_factory.get_trace', side_effect=lambda: trace_mock