Mocking Python Classes

Lately, I had to write unit tests using Pytest for a Python module. The module contains a class where other classes are initialize within its constructor.

As usual I created a fixture for this class to make it easy to write a test for each class method. At this point I ran into some issues when I tried to mock the different classes initiated in the constructor. The mocking didn't work, and instances of these classes were still being created.

After some research and combining a few different solutions I found online, I want to share how I managed to mock the classes.

Solution

Here is an example of the class I tried to mock:

class ClassA:
    def __init__(self):
        self.class_b = ClassB()
        self.class_c = ClassC()
        self.count = 0

We want to set a value for every field of this class during tests. This value can be None or a class mock, but we don't want initiations of the classes ClassB and ClassC.

In our case, let's decide that self.class_b and self.class_c should be mocks:

@pytest.fixture
def mock_class_b():
    class_b = Mock(spec=ClassB)
    return class_b

@pytest.fixture
def mock_class_c():
    class_c = Mock(spec=ClassC)
    return class_c

So a fixture for this class that serves our goal looks like this:

@pytest.fixture
def class_a_mock(mock_class_b, mock_class_c):
    with patch.object(target=ClassA, attribute="__init__", return_value=None) as mock_init:
        class_a = ClassA()
        class_a.class_b = mock_class_b
        class_a.class_c = mock_class_c
        class_b.count = 0
        return class_a

The important part is how to use the patch.object function, which is from unittest.mock module in Python. It is used in testing to temporarily replace an attribute of a given object wit a mock or another value.

Arguments

1. target=ClassA: the object (usually a class) whose attribute we want to patch.
2. attribute="__init__": the name of the attribute we want to patch.
3. return_value=None: replacing the __init__ method with a function that does nothing

In this way we can create mocked variables in our fixture.
Read more about patch.object

Testing Tips

I wrote this tutorial for this kind of cases where, for any reason, we cannot change the code of Class A. However, I generally recommend modifying the code if possible, not to change the logic, but to make it more testable.

Here are some examples of how to modify Class A to make it more testable:

Option 1: Pass instances of class B and class C as parameters.
This way, when we write the fixture, we can pass mocks instead of instances.

class ClassA:
    def __init__(self, class_b_instance, class_c_instance):
        self.class_b = class_b_instance
        self.class_c = class_c_instance
        self.count = 0

Option 2: Create a Boolean variable that indicates test mode.
This way we can decide which fields of class A will or will not get a value when it is initiated.

class ClassA:
    def __init__(self, test_mode=False):
        if not test_mode:
            self.class_b = ClassB()
            self.class_c = ClassC()
        self.count = 0

Option 3: Make class initiations in a separate method.
This approach gives us the choice to avoid calling set_class_variables in the test module.

class ClassA:
    def __init__(self):
        self.class_b = None
        self.class_c = None
        self.count = None

    def set_class_variables(self):
        self.class_b = ClassB()
        self.class_c = ClassC()
        self.count = 0

Hope this helps! :)