Python: mock environment variables with unittest
Update (2024-09-18): Updated this post with enterClassContext and enterContext from Python 3.11.
Sometimes, tests need to change environment variables. This is straightforward in tests using Python’s unittest, thanks to os.environ quacking like a dict, and the mock.patch.dict decorator/context manager.
(If you’re using pytest, see the pytest edition of this post.)
Adding environment variables
If you want to write a test that sets one or more environment variables, overriding existing values, you can use mock.patch.dict like this:
import os
from unittest import TestCase, mock
from example.settings import Colour, get_settings
class SettingsTests(TestCase):
@mock.patch.dict(os.environ, {"FROBNICATION_COLOUR": "ROUGE"})
def test_frobnication_colour(self):
colour = get_settings().frobnication_colour
self.assertEqual(colour, Colour.ROUGE)
You can apply this to all tests in a TestCase by applying it as a class decorator:
import os
from unittest import TestCase, mock
@mock.patch.dict(os.environ, {"FROBNICATION_COLOUR": "ROUGE"})
class SettingsTests(TestCase):
def test_frobnication_colour(self):
self.assertEqual(os.environ["FROBNICATION_COLOUR"], "ROUGE")
def test_frobnication_shade(self):
self.assertEqual(os.environ["FROBNICATION_COLOUR"], "ROUGE")
Note this wraps only methods starting test_, so setUp(), tearDown(), setUpClass(), etc. will use the unmocked environment. If you want to wrap the test case execution from start to end, you’ll want to wrap the mocker in setUpClass():
import os
from unittest import TestCase, mock
class SettingsTests(TestCase):
@classmethod
def setUpClass(cls):
cls.enterClassContext(
mock.patch.dict(os.environ, {"FROBNICATION_COLOUR": "ROUGE"})
)
# set from this point until tearDownClass()
def test_frobnication_colour(self):
self.assertEqual(os.environ["FROBNICATION_COLOUR"], "ROUGE")
The above example uses enterClassContext, new in Python 3.11, as I covered in a recent post. You can backport it to older versions for simpler tests.
Dynamic values
If you don’t know the keys or values you want to mock at import time, use the context manager form of mock.patch.dict within your test method:
import os
from unittest import TestCase, mock
from example.settings import Colour, get_settings
from tests.fixtures import get_mock_colour
class SettingsTests(TestCase):
def test_frobnication_colour(self):
with mock.patch.dict(os.environ, {"FROBNICATION_COLOUR": get_mock_colour()}):
colour = get_settings().frobnication_colour
self.assertEqual(colour, Colour.ROUGE)
Or, with enterContext in setUp():
import os
from unittest import TestCase, mock
from example.settings import Colour, get_settings
from tests.fixtures import get_mock_colour
class SettingsTests(TestCase):
def setUp(self):
self.enterContext(
mock.patch.dict(os.environ, {"FROBNICATION_COLOUR": get_mock_colour()})
)
# set from this point until tearDown()
def test_frobnication_colour(self):
colour = get_settings().frobnication_colour
self.assertEqual(colour, Colour.ROUGE)
Clearing before adding
If you want to clear everything from os.environ so only the given variables are set, you can do so by passing clear=True to mock.patch.dict:
import os
from unittest import TestCase, mock
from example.settings import get_settings
class SettingsTests(TestCase):
@mock.patch.dict(os.environ, {"FROBNICATION_COLOUR": "ROUGE"}, clear=True)
def test_frobnication_colour(self):
settings = get_settings()
self.assertEqual(settings.modified_names, {"FROBNICATION_COLOUR"})
Removing just a few variables
If you want to remove only a few variables, it gets a little more tricky. mock.patch.dict doesn’t have a way of removing select keys, so you need to build a dictionary of the keys to preserve and use that with clear=True:
import os
from unittest import TestCase, mock
from example.settings import get_settings
class SettingsTests(TestCase):
def test_frobnication_colour(self):
names_to_remove = {"FROBNICATION_COLOUR"}
modified_environ = {
k: v for k, v in os.environ.items() if k not in names_to_remove
}
with mock.patch.dict(os.environ, modified_environ, clear=True):
settings = get_settings()
self.assertEqual(settings.modified_names, set())
😸😸😸 Check out my new book on using GitHub effectively, Boost Your GitHub DX! 😸😸😸
One summary email a week, no spam, I pinky promise.
Related posts: