Add temp_env util function to temporarily set env variables with a context manager

Author: surister

tests/client/test_utils.py

@@ -0,0 +1,24 @@
+import os
+
+from tests.client.utils import temp_env
+
+
+def test_util_temp_env():
+    """
+    Verify that temp_env correctly injects and cleans up environment variables.
+    """
+    env = os.environ.copy()
+    new_env_var = {'some_var': '1'}
+
+    with temp_env(**new_env_var):
+        # Assert that the difference of the two environs is the new variable.
+        assert set(env.items()) ^ set(os.environ.items()) == set(new_env_var.items())
+
+    # key should now not exist in the current imported os.environ.
+    assert list(new_env_var.keys())[0] not in os.environ
+
+    import importlib
+    importlib.reload(os)
+
+    # key should still not exist even when reloading the os module.
+    assert list(new_env_var.keys())[0] not in os.environ

tests/client/utils.py

@@ -0,0 +1,25 @@
+import contextlib
+import os
+
+
+@contextlib.contextmanager
+def temp_env(**environment_variables):
+    """
+    Context manager that sets temporal environment variables within the scope.
+
+    Example:
+        with temp_env(some_key='value'):
+            import os
+            'some_key' in os.environ
+            # True
+        'some_key' in os.environ
+        # False
+    """
+    try:
+        for k, v in environment_variables.items():
+            os.environ[k] = str(v)
+            yield
+    finally:
+        for k, v in environment_variables.items():
+            os.unsetenv(k)
+            del os.environ[k]