How to verify Python mock calls with generator argument

user_1076480 0
Viewed 637

I have a function that takes an Iterable as an argument:

def print_me(vals: Iterable[str]) -> None:
    ...

It gets called with a generator:

def run():
    print_me((str(i) for i in range(10)))

I'd like to patch and mock print_me() and verify it gets called with specific parameters. With plain old lists, this is easy:

mock.assert_called_with(["0", "1", ...])

but with a generator, because it can't be matched on equality or replayed, this is trickier.

generator mocking unit-testing python
2 Answers

I found a few solutions. I can manually capture the argument, but it gets messier if the method is called more than once:

@patch('module.print_me')
def test(self, mock):
    captured = None
    mock.side_effect = lambda vals: captured = vals
    run()    
    mock.assert_called_once_with(mock.ANY)
    self.assertEqual(list(captured), ["0", "1", ...])

I can also proxy the mock. This makes assertions clean, but is a bit weird to think about, and the proxy() function needs to be updated every time the real function's arguments change.

@patch('module.print_me')
def test(self, proxy_mock):
    mock = mock.Mock()
    def proxy(vals):
        mock(list(vals))

    proxy_mock.side_effect = proxy

    run()    
    mock.assert_called_once_with(["0", "1"])
    mock.assert_called_once_with(["2", "3"])

There's also PyHamcrest. Hamcrest is really popular with Java tests, but doesn't seem to be as popular for Python. That said, it does what I want:

@patch('module.print_me')
def test(self, mock):
    run()
    mock.assert_called_once_with(match_equality(contains_exactly(["2", "3"])))
user_10764800

There is an easier way to test this, and that is by using the the call_args attribute of the mock object. call_args returns a tuple where the first item is the args and the second item is the kwargs. The syntax was slightly changed in Python 3.8 to allow for accessing these items by doing mock_object.call_args.args or mock_object.call_args.kwargs but I am running this on Python 3.7 so I am taking the more antiquated approach.

@patch("src.module.print_me")
def test_run(mock_print):
    run()
    args, _ = mock_print.call_args
    # grab the one and only call argument
    gen = args[0]
    expected = [str(i) for i in range(10)]
    
    assert list(gen) == expected

As you can see in the code snippet above, we grab the argument that was fed in to our mock object, consume it and then compare it against the expected value.

user_68178350
Related