How to patch a module level access of os.environ?

Question

I am writing test cases for a module module.py that imports from another module legacy.py. The legacy.py reads os.environ["some_var"] at module level. When I am trying to run test cases for module.py, they are failing with KeyError for some_var in os.environ.

This is how code looks like:

module.py

from legacy import db

def fun():
    pass

legacy.py

import os

env = os.environ["some_var"]

class db:
    def __init__(self):
        pass

When running test cases for module.py, I am getting KeyError: 'some_var'.

I tried patching os (at module level also putting it before importing from module.py in test file) but the import statement is run before it could be patched. I tried looking for similar question on StackOverflow but didn't find the exact problem I am facing here.

How can I mock it? Or any other suggestion to be able to run the test cases. Assume that I cannot modify the legacy.py file.

Answer 1

You could use the mock.patch.dict. There is an example in official doc.

Here is a fully functional example based in your question.

module.py

import os


def get_env_var_value():
    env_var = os.environ['PATH']
    return env_var

print(get_env_var_value())

test_module.py

from unittest.mock import patch
from unittest import main, TestCase

import module

class MyClassTestCase(TestCase):

    @patch.dict('os.environ', {'PATH': '/usr/sbin'})
    def test_get_env_var_value(self):
        self.assertEqual(module.get_env_var_value(), '/usr/sbin')

if __name__ == '__main__':
    main()

Doesn't matter if the environment var is loaded in outside or inside the class/method scope, because you will mock for all test scope.