Python mockito to stub an object for unit testing

mockito is used in the project where I recently joined. I checked the official site but it doesn’t have enough examples to understand the behavior.

So I tried to use it myself to understand how to use it. I hope it helps you too.

Definition of target functions

Firstly, I defined the following class. These functions’ behavior will be controlled under the test functions.

class DataChangeDetector:
    def trigger_callback(self, trigger, callback):
        if trigger:
            callback(1, 2)


class ForStub:
    def __init__(self, detector = DataChangeDetector()):
        self.__val = 1
        self.__detector = detector

    def func1(self, param1: int, param2: int) -> int:
        return param1 + param2

    def func1_name(self, *, param1: int, param2: int) -> int:
        return self.func1(param1, param2)

    def func2(self) -> None:
        return

    def func3(self) -> int:
        return self.__val

    def __private_func(self) -> str:
        return "private string value"

    def func4(self) -> str:
        return self.__private_func()

    def func5(self, trigger, callback) -> None:
        self.__detector.trigger_callback(trigger,callback)

Using when function

Stub for specific input

If you need to stub the function for a specific input, you can write it this way.

def test1_func1():
    instance = ForStub()
    when(instance).func1(1, 1).thenReturn(3)

    assert 3 == instance.func1(1, 1)
    with pytest.raises(invocation.InvocationError):
        instance.func1(1, 2)

Pass the target instance to when function followed by the function call with the specific arguments you want to stub. After that, you can set any value by thenReturn.

Even though the inputs are (1, 1), the function returns 3. If the arguments are different from the specified ones, it throws InvocationError.

If you need multiple inputs, you can just define it.

def test1_func1():
    instance = ForStub()
    when(instance).func1(1, 1).thenReturn(3)
    when(instance).func1(5, 1).thenReturn(99)

    assert 3 == instance.func1(1, 1)
    assert 99 == instance.func1(5, 1)

    with pytest.raises(invocation.InvocationError):
        instance.func1(1, 2)

Stub for arbitrary input

If you don’t need to specify specific inputs but want to return the dummy value for any input, you can use triple dots instead.

def test2_func1():
    instance = ForStub()
    when(instance).func1(...).thenReturn(3)

    assert 3 == instance.func1(1, 8)
    assert 3 == instance.func1(5, 5)

In this way, func1 always returns the same value for any inputs.

Stub for all instances

The previous way is to stub the function only for the specified instance. If you want to stub a function for all instances that can’t be controlled in the test function, you can pass the class name there instead.

def test3_func1():
    when(ForStub).func1(1, 1).thenReturn(3)

    instance = ForStub()
    assert 3 == instance.func1(1, 1)
    instance2 = ForStub()
    assert 3 == instance2.func1(1, 1)

Using when2 function

If you prefer, you can also use when2 function. The target function needs to be passed for this function.

def test4_func1():
    instance = ForStub()
    when2(instance.func1, ...).thenReturn(3)

    assert 3 == instance.func1(1, 19)
    assert 3 == instance.func1(5, 5)


def test5_func1():
    instance = ForStub()
    when2(instance.func1, 1, 1).thenReturn(3)

    assert 3 == instance.func1(1, 1)
    with pytest.raises(invocation.InvocationError):
        instance.func1(9, 9)

Return None

If you return None, you don’t need to pass any argument to thenReturn.

def test1_func2():
    instance = ForStub()
    when(instance).func2().thenReturn()

    assert None == instance.func2()

Return different result depending on the call count

If it’s necessary to control the return value depending on the call count, you can just add the return value to thenReturn.

def test2_func2():
    instance = ForStub()
    when(instance).func2().thenReturn(1, 2, 3, 4)

    assert 1 == instance.func2()
    assert 2 == instance.func2()
    assert 3 == instance.func2()
    assert 4 == instance.func2()

List value can be set too.

def test3_func2():
    instance = ForStub()
    when(instance).func2().thenReturn([1, 2], (3, 4), [5, 6])

    assert [1, 2] == instance.func2()
    assert (3, 4) == instance.func2()
    assert [5, 6] == instance.func2()

You can also write it in dot chaining.

def test4_func2():
    instance = ForStub()
    when(instance).func2().thenReturn(1).thenReturn(5).thenReturn(8)

    assert 1 == instance.func2()
    assert 5 == instance.func2()
    assert 8 == instance.func2()

Stub private member value

Sometimes, a private member needs to be controlled for various reasons although this is a kind of coding smell.

If you try to change the value in the following ways, it throws Attribute Error.

def test1_func3():
    instance = ForStub()
    with pytest.raises(AttributeError):
        instance.__val == 9876

    with pytest.raises(AttributeError):
        when2(instance.__val).thenReturn(33)

    assert 1 == instance.func3()

The private member can’t be accessed from outside of the class. So we need to change the value in a different way. This is the answer.

def test2_func3():
    instance = ForStub()
    # {'_ForStub__val': 1}
    print(instance.__dict__)
    instance._ForStub__val = 33
    assert 33 == instance.func3()

You can change the value in this way. class_instance._ClassName__private_variable_name.

Stub private member function

Likewise, a private function can be stubbed in the same way.

def test1_func4():
    instance = ForStub()
    # private string value
    print(instance.func4())

    def stub_private_func():
        return "dummy value"

    instance._ForStub__private_func = stub_private_func
    assert "dummy value" == instance.func4()

How to replace a method to invoke a callback with thenAnswer

The function might have a callback. I added trigger parameter to make it easy. The actual code is often more complex. Image that DataChangeDetector is not our module. In this case, we need to change the behavior to invoke the callback.

class DataChangeDetector:
    def trigger_callback(self, trigger, callback):
        if trigger:
            callback(1, 2)

class ForStub:
    def __init__(self, detector = DataChangeDetector()):
        self.__val = 1
        self.__detector = detector

    def func5(self, trigger, callback) -> None:
        self.__detector.trigger_callback(trigger,callback)

Let’s check how it is used first. The callback is called only if trigger is set to true.

def test1_func5():
    instance = ForStub()
    value = [0]

    def callback(a, b):
        value[0] = a + b

    instance.func5(True, callback)
    assert value[0] == 3


def test2_func5():
    instance = ForStub()
    value = [0]

    def callback(a, b):
        value[0] = a + b

    instance.func5(False, callback)
    assert value[0] == 0

If the trigger is false, the value keeps 0.

To replace a method, we can use thenAnswer. The instance of DataChangeDetector needs to be stubbed. So, it needs to be injected in the test.

def test3_func5():
    detector = DataChangeDetector()
    instance = ForStub(detector)
    value = [0]

    def callback(a, b):
        value[0] = a + b

    when(detector).trigger_callback(...).thenAnswer(lambda a, b: callback(5, 5))
    instance.func5(False, callback)
    assert value[0] == 10

lambda is used here but it can be replaceable with a normal function. The trigger is set to false but the callback is called with arbitrary parameters.

How to invoke a callback for a internally defined callback

It’s better to show another code in a practical way. Let’s say we have subscribe method to know the data change. We have update method that notifies the data change to the listener.

class DataChangeDetector:
    _current_value: Optional[int] = None
    _listener: Optional[Callable[[int], None]] = None

    def subscribe(self, callback: Callable[[int], None]):
        self._listener = callback

    def update(self, value: int):
        if self._current_value != value:
            self._current_value = value
            if self._listener is not None:
                self._listener(value)

class ForStub:
    # ... other functions ...

    def subscribe(self, callback: Callable[[int], None]) -> None:
        def on_updated(new_value: int):
            if new_value < 10:
                callback(5)
            else:
                callback(new_value)

        self.__detector.subscribe(on_updated)

If we know that update method notifies an error via the listener, we can call it in the unit test in the following way.

def test1_subscribe():
    detector = DataChangeDetector()
    instance = ForStub(detector)

    value = [0]

    def callback(new_value):
        value[0] = new_value

    instance.subscribe(callback)

    detector.update(9)
    assert value[0] == 5
    detector.update(10)
    assert value[0] == 10

However, it’s not always so simple. The update method might not be exposed to the user. The data might be updated by someone else. For example, if it’s a digital thermometer, the device will trigger the callback.

In this case, we can write it in the following way.

def test2_subscribe():
    detector = DataChangeDetector()
    instance = ForStub(detector)

    captor = matchers.captor()
    when(detector).subscribe(captor)

    value = [0]

    def callback(new_value):
        value[0] = new_value

    instance.subscribe(callback)

    on_updated = captor.value

    on_updated(9)
    assert value[0] == 5
    on_updated(10)
    assert value[0] == 10

The key point is to use matchers.captor() to capture the parameter. It can store the parameters that are specified in the method. It stores the callback specified in self.__detector.subscribe. Since its value is stored in value property, we can call the callback in this way. We don’t have to consider the logic in update method.

How to verify the received parameters with spy

If you want to check the received parameters, you can use spy or spy2. I use spy2 here. Once the target method is specified in spy2, it can be verified with verify function.

def test1_verify_with_spy2():
    instance = ForStub()
    spy2(instance.func1)
    instance.func1(1, 1)
    instance.func1(2, 2)
    verify(instance, atleast=1).func1(...)
    verify(instance, times=2).func1(...)
    verify(instance, atleast=1).func1(1, 1)
    verify(instance, atleast=1).func1(2, 2)

    # Error
    # verify(instance, times=2).func1(1, 1)

If you just want to check if the method is called, just call verify with the method. It’s the first one. If you need to check the exact call count, set the value to times. You can check if the method is called with the exact parameters by setting them to the function call.

Note that parameter name must be added if the method uses it.

def func1_name(self, *, param1: int, param2: int) -> int:
    return self.func1(param1, param2)

def test2_verify_with_spy2():
    instance = ForStub()
    spy2(instance.func1_name)
    instance.func1_name(param1=1, param2=1)
    instance.func1_name(param1=2, param2=2)
    verify(instance, times=2).func1_name(...)
    verify(instance, atleast=1).func1_name(param1=1, param2=1)
    verify(instance, atleast=1).func1_name(param1=2, param2=2)

How to verify the received parameters with when

Sometimes the target method needs to be stubbed too. It’s possible to stub it followed by checking the received parameters. Once the method is stubbed, the same features of spy can be used.

def test_verify_with_when():
    instance = ForStub()
    when(instance).func1(...).thenReturn(5)
    instance.func1(1, 1)
    instance.func1(2, 2)
    verify(instance, atleast=1).func1(...)
    verify(instance, atleast=1).func1(1, 1)
    verify(instance, atleast=1).func1(2, 2)
    verify(instance, times=2).func1(...)