set env in config instead of using os to do it

Author: duncanista

datadog_lambda/cold_start.py

@@ -1,4 +1,3 @@
-import os
 import time
 from typing import List, Hashable
 import logging
@@ -74,9 +73,7 @@ def is_managed_instances_mode():
     Returns:
         bool: True if running in managed instances mode, False otherwise
     """
-    return (
-        os.environ.get("AWS_LAMBDA_INITIALIZATION_TYPE") == "lambda-managed-instances"
-    )
+    return config.aws_lambda_initialization_type == "lambda-managed-instances"
 
 
 class ImportNode(object):

datadog_lambda/config.py

@@ -98,6 +98,7 @@ def _resolve_env(self, key, default=None, cast=None, depends_on_tracing=False):
     integration_test = _get_env("DD_INTEGRATION_TEST", "false", as_bool)
 
     aws_lambda_function_name = _get_env("AWS_LAMBDA_FUNCTION_NAME")
+    aws_lambda_initialization_type = _get_env("AWS_LAMBDA_INITIALIZATION_TYPE")
 
     @property
     def function_name(self):