Temporarily modify the current process's environment

user_55246 0
Viewed 39673

I use the following code to temporarily modify environment variables. 

@contextmanager
def _setenv(**mapping):
    """``with`` context to temporarily modify the environment variables"""
    backup_values = {}
    backup_remove = set()
    for key, value in mapping.items():
        if key in os.environ:
            backup_values[key] = os.environ[key]
        else:
            backup_remove.add(key)
        os.environ[key] = value

    try:
        yield
    finally:
        # restore old environment
        for k, v in backup_values.items():
            os.environ[k] = v
        for k in backup_remove:
            del os.environ[k]

This with context is mainly used in test cases. For example,

def test_myapp_respects_this_envvar():
    with _setenv(MYAPP_PLUGINS_DIR='testsandbox/plugins'):
        myapp.plugins.register()
        [...]

My question: is there a simple/elegant way to write _setenv? I thought about actually doing backup = os.environ.copy() and then os.environ = backup .. but I am not sure if that would affect the program behavior (eg: if os.environ is referenced elsewhere in the Python interpreter).

environment-variables python
5 Answers
_environ = dict(os.environ)  # or os.environ.copy()
try:

    ...

finally:
    os.environ.clear()
    os.environ.update(_environ)
Ignacio_Vazquez-Abrams_208620

I was looking to do the same thing but for unit testing, here is how I have done it using the unittest.mock.patch function:

def test_function_with_different_env_variable():
    with mock.patch.dict('os.environ', {'hello': 'world'}, clear=True):
        self.assertEqual(os.environ.get('hello'), 'world')
        self.assertEqual(len(os.environ), 1)

Basically using unittest.mock.patch.dict with clear=True, we are making os.environ as a dictionary containing solely {'hello': 'world'}.

  • Removing the clear=True will let the original os.environ and add/replace the specified key/value pair inside {'hello': 'world'}.

  • Removing {'hello': 'world'} will just create an empty dictionary, os.envrion will thus be empty within the with.

user_77479420

In pytest you can temporarily set an environment variable using the monkeypatch fixture. See the docs for details. I've copied a snippet here for your convenience.

import os
import pytest
from typing import Any, NewType

# Alias for the ``type`` of monkeypatch fixture.
MonkeyPatchFixture = NewType("MonkeyPatchFixture", Any)


# This is the function we will test below to demonstrate the ``monkeypatch`` fixture.
def get_lowercase_env_var(env_var_name: str) -> str:
    """
    Return the value of an environment variable. Variable value is made all lowercase.

    :param env_var_name:
        The name of the environment variable to return.
    :return:
        The value of the environment variable, with all letters in lowercase.
    """
    env_variable_value = os.environ[env_var_name]
    lowercase_env_variable = env_variable_value.lower()
    return lowercase_env_variable


def test_get_lowercase_env_var(monkeypatch: MonkeyPatchFixture) -> None:
    """
    Test that the function under test indeed returns the lowercase-ified
    form of ENV_VAR_UNDER_TEST.
    """
    name_of_env_var_under_test = "ENV_VAR_UNDER_TEST"
    env_var_value_under_test = "EnvVarValue"
    expected_result = "envvarvalue"
    # KeyError because``ENV_VAR_UNDER_TEST`` was looked up in the os.environ dictionary before its value was set by ``monkeypatch``.
    with pytest.raises(KeyError):
        assert get_lowercase_env_var(name_of_env_var_under_test) == expected_result
    # Temporarily set the environment variable's value.
    monkeypatch.setenv(name_of_env_var_under_test, env_var_value_under_test)
    assert get_lowercase_env_var(name_of_env_var_under_test) == expected_result


def test_get_lowercase_env_var_fails(monkeypatch: MonkeyPatchFixture) -> None:
    """
    This demonstrates that ENV_VAR_UNDER_TEST is reset in every test function.
    """
    env_var_name_under_test = "ENV_VAR_UNDER_TEST"
    expected_result = "envvarvalue"
    with pytest.raises(KeyError):
        assert get_lowercase_env_var(env_var_name_under_test) == expected_result
user_55894390
Related