How to Mock Environment Variables in pytest

2020-10-13 Our environment - planet earth

Sometimes tests need to change environment variables. This is fairly straightforward in pytest, thanks to os.environ quacking like a dict, and the mock.patch.dict decorator/context manager.

(If you’re not using pytest, or use TestCase classes with pytest, see the unittest edition of this post.)

`mock.patch` or `monkeypatch` ?

Update (2020-10-15): Added this section, thanks to Tom Grainger on Twitter for the hint about monkeypatch.

pytest comes with a monkeypatch fixture which does some of the same things as mock.patch. This post uses mock.patch, since it’s a more powerful and general purpose tool. But you might prefer monkeypatch - check out the monkeypatch documentation for environment variables.

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 mock

from example.settings import Colour, get_settings


@mock.patch.dict(os.environ, {"FROBNICATION_COLOUR": "ROUGE"})
def test_frobnication_colour():
    colour = get_settings().frobnication_colour
    assert colour == Colour.ROUGE

You can apply this to all tests in a module by creating a local auto-used pytest fixture that uses mock.patch.dict:

import os
from unittest import mock

import pytest


@pytest.fixture(autouse=True)
def mock_settings_env_vars():
    with mock.patch.dict(os.environ, {"FROBNICATION_COLOUR": "ROUGE"}):
        yield


def test_frobnication_colour():
    assert os.environ["FROBNICATION_COLOUR"] == "ROUGE"


def test_frobnication_shade():
    assert 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 function:

import os
from unittest import mock

from example.settings import Colour, get_settings
from tests.fixtures import get_mock_colour


def test_frobnication_colour():
    with mock.patch.dict(os.environ, {"FROBNICATION_COLOUR": get_mock_colour()}):
        colour = get_settings().frobnication_colour

    assert 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 mock

from example.settings import get_settings


@mock.patch.dict(os.environ, {"FROBNICATION_COLOUR": "ROUGE"}, clear=True)
def test_frobnication_colour():
    settings = get_settings()
    assert 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 mock

from example.settings import get_settings


def test_frobnication_colour():
    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()
    assert 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 Python’s unittest

Tags: python