Mocking Functions with Type Hints in Python

“I want to add type hints to test code as well, without exception.”

With that thought in mind while writing Python, I got curious about whether I could create a mock for a function that has type hints. I managed to create something that seems to work reasonably well, so I'd like to introduce it.

The idea of what I want to achieve looks like this:

function_mock = FunctionMock(args).returns(return_value)
ret = function_mock(value) # Assertion is performed here
use(ret) # The specified return value is returned

Mock Implementation Code

from __future__ import annotations

from typing import Callable, Generic, List, Optional, ParamSpec, TypeVar


Params = ParamSpec('Params')
Return = TypeVar('Return')


class Call(Generic[Params, Return]):
    retval: Optional[Return]

    def __init__(self, *args: Params.args, **kwargs: Params.kwargs) -> None:
        self.args = args
        self.kwargs = kwargs
        self.retval = None

    def returns(self, retval: Return) -> Call:
        self.retval = retval
        return self


class FunctionMock(Callable[Params, Return]):
    calls: List[Call[Params, Return]]
    count: int

    def __init__(self, calls: List[Call[Params, Return]]) -> None:
        self.calls = calls
        self.count = 0

    def __call__(self, *args: Params.args, **kwargs: Params.kwargs) -> Return:
        assert self.count < len(self.calls), f'Too many calls to function: want {len(self.calls)}.'
        call = self.calls[self.count]
        assert args == call.args, f'Argument mismatch: want {call.args}, got {args}'
        assert kwargs == call.kwargs, f'Keyword argument mismatch: want {call.kwargs}, got {kwargs}'
        self.count += 1
        return call.retval

    def close(self):
        assert self.count == len(self.calls), f'Not enough calls to function: want {len(self.calls)}, got {self.count}.'

    def __enter__(self):
        return self

    def __exit__(self, ex_type, ex_value, traceback):
        self.close()
        return ex_type is None

Explanation

A key point when adding type hints to a function mock is that "how to handle argument type hints" is surprisingly tricky. This is because, in Python, you cannot simply compare argument sequences as tuples due to the existence of keyword arguments. Therefore, we will use ParamSpec.

ParamSpec is a class for creating special type parameters to represent function arguments. By defining a type parameter like P = ParamSpec('P'), you can define a function type as Callable[P, R], and you can add type hints to arguments and keyword arguments using the types P.args and P.kwargs. However, there is one constraint: P.args and P.kwargs can only be used as type hints for *args and **kwargs, respectively. In other words, they can practically only be used for the type hints of a function's arguments. Therefore, if you want to use these type hints, you need to find a way to express them skillfully using functions. I will proceed while keeping this point in mind.

Defining Type Parameters

To specify the types of arguments and return values when creating a mock, the argument types are represented with ParamSpec, and the return value type is represented with TypeVar.

Params = ParamSpec('Params')
Return = TypeVar('Return')

Objectifying Function Calls

The behavior required for a function mock consists of two things: "matching the arguments when called" and "returning a value in place of the hidden behavior." As mentioned earlier, since ParamSpec's args and kwargs can only be used in function argument type hints, we use them in the class constructor. Ideally, I wanted to be able to specify the return value in the constructor as well, but since that felt slightly off from the Callable[Params, Return] typing, I added a separate returns method (if there's a better way to avoid this, please let me know).

class Call(Generic[Params, Return]):
    retval: Optional[Return]

    def __init__(self, *args: Params.args, **kwargs: Params.kwargs) -> None:
        self.args = args
        self.kwargs = kwargs
        self.retval = None

    def returns(self, retval: Return) -> Call:
        self.retval = retval
        return self

Mock Body

For the mock body, three behaviors are expected: "being given a list of arguments to match and values to return when called," "verifying that the expected arguments are passed and returning the specified value each time it is called," and "checking whether it was called the specified number of times." FunctionMock implements these straightforwardly.

class FunctionMock(Callable[Params, Return]):
    calls: List[Call[Params, Return]]
    count: int

    def __init__(self, calls: List[Call[Params, Return]]) -> None:
        self.calls = calls
        self.count = 0

    def __call__(self, *args: Params.args, **kwargs: Params.kwargs) -> Return:
        assert self.count < len(self.calls), f'Too many calls to function: want {len(self.calls)}.'
        call = self.calls[self.count]
        assert args == call.args, f'Argument mismatch: want {call.args}, got {args}'
        assert kwargs == call.kwargs, f'Keyword argument mismatch: want {call.kwargs}, got {kwargs}'
        self.count += 1
        return call.retval

    def close(self):
        assert self.count == len(self.calls), f'Not enough calls to function: want {len(self.calls)}, got {self.count}.'

    def __enter__(self):
        return self

    def __exit__(self, ex_type, ex_value, traceback):
        self.close()
        return ex_type is None

Testing the Mock

To demonstrate how to use it, let's write a test with pytest to confirm that the mock works as expected.

def test_function_mock():
    calls = [Call[[int], int](i).returns(i * 2) for i in range(10)]
    with FunctionMock[[int], int](calls) as double_mock:
        for i in range(10):
            double_mock(i)

Conclusion

Function mocks protected by type hints are possible! While I feel I could create a feature-rich mock like unittest.Mock if I really wanted to, as a minimalist, I think this is just about right. I don't have any plans to package it (as of this writing), but depending on the response, I might release it as a small package. Happy type-hinting life.