Python mock os.environ used inside a class

Question

I'm trying to mock os.environ inside a class but I just can't get it right. Here's my structure:

#file.py
import os

class MyClass():
    connection_url = os.environ['DB']

#some code

And here's my test (latest try, anyways):

#test.py
from unittest import TestCase
from unittest.mock import patch
from file import MyClass

class TestMyClass(TestCase):
    @patch.dict('file.os.environ', {'DB' : 'Dummy' })
    def setUp(self):
         self.class = MyClass()

#some testing

This is failing miserably, raising KeyError 'DB'... Can someone help me? I'm new to python unittesting. I researched some blogs and stackoverflow, tried some solutions but couldn't get it right.

Thanks in advance!

Answer 1

The problem here is that connection_url is set when the class is created (at import time). If you want to mock it, you can patch the attribute on the class itself before you use that attribute:

class TestMyClass(TestCase):
    @patch.object(file.MyClass, 'connection_url', 'Dummy')
    def setUp(self):
        self.instance = MyClass()

You'll still have to deal with the potential KeyError at import time -- i.e. you'll need to make sure that your test runner has a dummy value in the environment before you import file or else you'll have to modify the code in file.py to make sure that it doesn't raise a KeyError. There are a few strategies here. You can just suppress the KeyError letting connection_url = default_value if it wasn't set in the environment:

class MyClass():
    connection_url = os.environ.get('DB', default_value)

Or you can move the connection_url fetching into __init__:

class MyClass():
    def __init__(self):
        connection_url = os.environ['DB']

This latter code has the additional benefit that your patch.dict should start to work.

Note that if you choose the either of the first two approaches, the instance will only have the patched connection_url during the setUp method -- The subsequent tests won't have that patched. It's quite likely that this will be a show-stopper and you'll need to create a new instance in each test that you have:

class TestMyClass(TestCase):
    @patch.object(file.MyClass, 'connection_url', 'Dummy')
    def test_first_thing(self):
        self.instance = MyClass()
        ...

    @patch.object(file.MyClass, 'connection_url', 'Dummy')
    def test_second_thing(self):
        self.instance = MyClass()
        ...