Consolidate and fix Environ fixtures

Author: TimPansino

tests/agent_unittests/test_utilization_settings.py

@@ -17,6 +17,7 @@
 from importlib import reload
 
 import pytest
+from testing_support.fixtures import Environ
 
 # these will be reloaded for each test
 import newrelic.config
@@ -65,25 +66,9 @@
     "NEW_RELIC_HEROKU_DYNO_NAME_PREFIXES_TO_SHORTEN": "meow wruff",
 }
 
-INITIAL_ENV = os.environ
 
 # Tests for loading settings and testing for values precedence
 
-
-class Environ:
-    def __init__(self, env_dict):
-        self.env_dict = {}
-        for key in env_dict.keys():
-            self.env_dict[key] = str(env_dict[key])
-
-    def __enter__(self):
-        os.environ.update(self.env_dict)
-
-    def __exit__(self, *args, **kwargs):
-        os.environ.clear()
-        os.environ[:] = INITIAL_ENV[:]
-
-
 def reset_agent_config(ini_contents, env_dict):
     @function_wrapper
     def reset(wrapped, instance, args, kwargs):

tests/cross_agent/test_pcf_utilization_data.py

@@ -16,12 +16,12 @@
 import os
 
 import pytest
+from testing_support.fixtures import Environ
 from testing_support.validators.validate_internal_metrics import validate_internal_metrics
 
 from newrelic.common.utilization import PCFUtilization
 
 CURRENT_DIR = os.path.dirname(os.path.realpath(__file__))
-INITIAL_ENV = os.environ
 FIXTURE = os.path.normpath(os.path.join(CURRENT_DIR, "fixtures", "utilization_vendor_specific", "pcf.json"))
 
 _parameters_list = ["testname", "env_vars", "expected_vendors_hash", "expected_metrics"]
@@ -42,27 +42,6 @@ def _parametrize_test(test):
 _pcf_tests = [_parametrize_test(t) for t in _load_tests()]
 
 
-class Environ:
-    def __init__(self, env_dict):
-        env_dict = env_dict or {}
-        cleaned_env_dict = {}
-        for key, val in env_dict.items():
-            if val is None:
-                continue
-            elif not isinstance(val, str):
-                cleaned_env_dict[key] = val.encode("utf-8")
-            else:
-                cleaned_env_dict[key] = val
-        self.env_dict = cleaned_env_dict
-
-    def __enter__(self):
-        os.environ.update(self.env_dict)
-
-    def __exit__(self, *args, **kwargs):
-        os.environ.clear()
-        os.environ[:] = INITIAL_ENV[:]
-
-
 class MockResponse:
     def __init__(self, code, body):
         self.code = code

tests/testing_support/fixtures.py

@@ -1335,17 +1335,29 @@ def check_error_attributes(
 class Environ:
     """Context manager for setting environment variables temporarily."""
 
-    def __init__(self, **kwargs):
-        self._original_environ = os.environ
-        self._environ_dict = kwargs
+    def __init__(self, env_dict=None, **kwargs):
+        self._environ_dict = {}
+
+        env_dict = env_dict or kwargs
+        for key, val in env_dict.items():
+            if val is None:
+                continue
+            elif isinstance(val, str):
+                self._environ_dict[key] = val
+            elif isinstance(val, bytes):
+                self._environ_dict[key] = val.encode("utf-8")
+            else:
+                self._environ_dict[key] = str(val)
 
     def __enter__(self):
+        self._original_environ = os.environ.copy()
+
         for key, val in self._environ_dict.items():
             os.environ[key] = str(val)
 
     def __exit__(self, exc, val, tb):
         os.environ.clear()
-        os.environ[:] = self._original_environ[:]
+        os.environ.update(self._original_environ)
 
 
 class TerminatingPopen(subprocess.Popen):