How to Mock Environment Variables in Python’s unittest

2020-10-13 Our environment - mother gaia

Sometimes tests need to change environment variables. This is fairly 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 create and start the mocker in setUpClass(), and stop it tearDownClass():

import os
from unittest import TestCase, mock


class SettingsTests(TestCase):
    @classmethod
    def setUpClass(cls):
        cls.env_patcher = mock.patch.dict(os.environ, {"FROBNICATION_COLOUR": "ROUGE"})
        cls.env_patcher.start()

        super().setUpClass()

    @classmethod
    def tearDownClass(cls):
        super().tearDownClass()

        cls.env_patcher.stop()

    def setUp(self):
        super().setUp()
        self.assertEqual(os.environ["FROBNICATION_COLOUR"], "ROUGE")

    def test_frobnication_colour(self):
        self.assertEqual(os.environ["FROBNICATION_COLOUR"], "ROUGE")

Dynamic Values

If you don’t know the keys or values you want to mock at import time, you’ll need to 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)

Clearing

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

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())

Fin

I hope this helps you with your testing journey,

—Adam

Learn how to make your tests run quickly in my book Speed Up Your Django Tests.

One summary email a week, no spam, I pinky promise.

Related posts:

How to Mock Environment Variables in pytest

Tags: python