Move flush timeout calculation into its own function

Author: jbfenton

instrumentation/opentelemetry-instrumentation-aws-lambda/src/opentelemetry/instrumentation/aws_lambda/__init__.py

@@ -433,6 +433,31 @@ def _instrumented_lambda_handler_call( # noqa pylint: disable=too-many-branches
  )
 
 
+def determine_flush_timeout() -> int:
+ """
+ Determine the flush timeout for the Lambda function.
+
+ Returns: Flush timeout in milliseconds
+ """
+
+ flush_timeout_env = os.environ.get(
+ OTEL_INSTRUMENTATION_AWS_LAMBDA_FLUSH_TIMEOUT, None
+ )
+
+ flush_timeout = 30000
+
+ try:
+ if flush_timeout_env is not None:
+ flush_timeout = int(flush_timeout_env)
+ except ValueError:
+ logger.warning(
+ "Could not convert OTEL_INSTRUMENTATION_AWS_LAMBDA_FLUSH_TIMEOUT value %s to int",
+ flush_timeout_env,
+ )
+
+ return flush_timeout
+
+
 class AwsLambdaInstrumentor(BaseInstrumentor):
  def instrumentation_dependencies(self) -> Collection[str]:
  return _instruments
@@ -462,19 +487,6 @@ def _instrument(self, **kwargs):
  self._wrapped_function_name,
  ) = lambda_handler.rsplit(".", 1)
 
- flush_timeout_env = os.environ.get(
- OTEL_INSTRUMENTATION_AWS_LAMBDA_FLUSH_TIMEOUT, None
- )
- flush_timeout = 30000
- try:
- if flush_timeout_env is not None:
- flush_timeout = int(flush_timeout_env)
- except ValueError:
- logger.warning(
- "Could not convert OTEL_INSTRUMENTATION_AWS_LAMBDA_FLUSH_TIMEOUT value %s to int",
- flush_timeout_env,
- )
-
  disable_aws_context_propagation = kwargs.get(
  "disable_aws_context_propagation", False
  ) or os.getenv(
@@ -488,7 +500,7 @@ def _instrument(self, **kwargs):
  _instrument(
  self._wrapped_module_name,
  self._wrapped_function_name,
- flush_timeout,
+ flush_timeout=determine_flush_timeout(),
  event_context_extractor=kwargs.get(
  "event_context_extractor", _default_event_context_extractor
  ),

instrumentation/opentelemetry-instrumentation-aws-lambda/tests/test_util_functions.py

@@ -0,0 +1,35 @@
+# Copyright The OpenTelemetry Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from opentelemetry.instrumentation.aws_lambda import determine_flush_timeout
+from opentelemetry.instrumentation.aws_lambda import (
+ OTEL_INSTRUMENTATION_AWS_LAMBDA_FLUSH_TIMEOUT,
+)
+from unittest.mock import patch
+
+
+class TestDetermineFlushTimeout:
+ default = 30000
+ timeout_var = OTEL_INSTRUMENTATION_AWS_LAMBDA_FLUSH_TIMEOUT
+
+ def test_flush_timeout_default(self):
+ assert determine_flush_timeout() == self.default
+
+ @patch.dict("os.environ", {timeout_var: "1000"})
+ def test_flush_timeout_env_var(self):
+ assert determine_flush_timeout() == 1000
+
+ @patch.dict("os.environ", {timeout_var: "abc"})
+ def test_flush_timeout_env_var_invalid(self):
+ assert determine_flush_timeout() == self.default