How to instruct a magic mock on how it should treat its arguments

Question

I've run into the following (edge?) case that I don't know how to handle properly. The general problem is that

• I have a function that I want to test
• in that function I call an external function with a generator comprehension as its argument
• in my tests, I mock the external function away
• now the prod code and the tested code differ: in prod, the generator is consumed, the mock doesn't do that

---

Here is a reduced example of what it looks like in my codebase:

import itertools
import random


def my_side_effects():
    # imaginge itertools.accumulate was some expensive strange function
    # that consumes an iterable
    itertools.accumulate(random.randint(1, 5) for _ in range(10))


def test_my_side_effects(mocker):
    my_mocked_func = mocker.patch('itertools.accumulate')

    my_side_effects()

    # make sure that side-effects took place. can't do much else.
    assert my_mocked_func.call_count == 1

The test runs just fine, and is good enough for all I care. But when I run coverage on the code, the situation that I described in the abstract becomes apparent:

----------- coverage: platform linux, python 3.8.0-final-0 -----------
Name                                   Stmts   Miss Branch BrPart  Cover   Missing
----------------------------------------------------------------------------------
[...]
my_test_case.py                            5      0      2      1    86%   6->exit
[...]
----------------------------------------------------------------------------------
# something like this, the ->exit part on the external call is the relevant part

Explanation of the ->exit syntax in coverage.py. Given that the comprehension could execute relevant business logic that I actually do want to run, the missed coverage is relevant. It's just calling random.randint here, but it could do anything.

---

Workarounds:

1. I can just use a list comprehension instead. The code is called and everybody is happy. Except me, who has to modify their backend in order to mollify tests.
2. I can reach into the mock during the test, grab the call arg, and unroll it by hand. This will probably look godawful.
3. I can monkeypatch the function instead of using a magicmock, something like monkeypatch.setattr('itertools.accumulate', lambda x: [*x]) would be quite descriptive. But I would lose the ability to make call assertions like in my example.

What I would consider a good solution would be something like this, which sadly doesn't exist:

def test_my_side_effects(mocker):
    my_mocked_func = mocker.patch('itertools.accumulate')

    # could also take "await", and assign treatments by keyword
    my_mocked_func.arg_treatment('unroll')  

    my_side_effects()

    # make sure that side-effects took place. can't do much else.
    assert my_mocked_func.call_count == 1

Answer 1

You are correct that there is missing coverage here: in fact, since the accumulate was never consumed, you could even have:

itertools.accumulate(ERRORERRORERROR for _ in range(10))

And your existing test would still pass (the obvious error just got mocked away).

To address this, use the side_effect of the mock:

my_mocked_func = mocker.patch('itertools.accumulate', side_effect=list)

When using a callable as a mock's side_effect, it gets called with the same arguments as the mock, and the return value of this callable is used as the return value of the mock (note: that means you may also assert on the returned value here rather than just the blunt call_count assertion).

That will allow you to consume the generator and get 100% coverage here.

Answer 2

Doing it the old way:

import itertools

def func():
    return list(itertools.izip(["a", "b", "c"], [1, 2, 3]))

def test_mock():
    callargs = []
    def mock_zip(*args):
        callargs.append(args)
        for arg in args:
            list(arg)
        yield ("a", 1)
        yield ("b", 2)

    old_izip = itertools.izip
    itertools.izip = mock_zip

    result = func()

    itertools.izip = old_izip

    assert 1 == len(callargs), "oops, not called once"
    assert result == [("a", 1), ("b", 2)], "oops, wrong result"

    print("success")