How to assert that a mock function is called with a generator?

Question

I'm using mock with py.test for unittests.

One function under test looks like the one below:

def convert(rows):
    rows = (preprocess(r) for r in rows)
    return batch_process(True, rows)

In the test, I would mock preprocess and batch_process . For preprocess , I use side_effect to configure multiple return values, one for each of the rows.

How do I assert that the second argument passed to batch_process is a generator expression of preprocess ed rows?

Answer 1

Consider to simplify your test: what you need to test is that the second argument when is traversed as a generator contains the values returned by preprocess() for each row.

I will just sketch how I can eventually write it :

from mock import Mock, call

preprocess_sentinels = [Mock() for _ in range(5)]
mock_preprocess.side_effect = preprocess_sentinels
rows = [Mock() for _ in range(5)]

convert(rows)

mock_preprocess.assert_has_calls([call(r) for r in rows])
args, kwargs = mock_batch_process.call_args
assert args[0]
for v, e in zip(args[1], preprocess_sentinels):
    assert v == e

Sure, in that way you don't check strictly it is a generator, but you check convert() behavior. Check exactly how you have implemented it (by generator or by list) is not what a test should do, it is an implementation detail and not a behavior: you should write the best test that works in both cases.

Answer 2

I've found a way to check this, but it is little verbose, but works for me in my project

from mock import Mock, patch
from types import GeneratorType

def test():
    mock_batch_process = patch('module.calling.batch_process')
    mock_preprocess = patch('module.calling.preprocess')
    # our test data for rows
    rows = ...

    convert(rows)

    assert mock_batch_process.call_count == 1
    args, kwargs = mock_batch_process.call_args
    # get an arg that is expected to be a generator, second arg in this case
    preprocess_rows = args[1]
    # check that our call arg is a generator
    assert isinstance(preprocess_rows, GeneratorType)
    # check that it contains all expected values
    assert list(preprocess_rows) == [mock_preprocess(row) for row in rows]