Throw exception after first call

Question

I have a loop where I handle adding records to a zip file. I have mocked my zipfile object and I want to raise a exception to verify that my logic for handling large zipfiles will work correctly.

is there a way with MagicMocks or plain mocks to accept the first call, but raise an exception on the second?

Answer 1

The simplest way is to use side_effect that accept either iterable, callable or Exception (class or instance)

Alternatively side_effect can be an exception class or instance. In this case the exception will be raised when the mock is called.

As showed in Quick Guide you can use side_effect to raise an exception simply by

>>> mock = Mock(side_effect=KeyError('foo'))
>>> mock()
Traceback (most recent call last):
 ...
KeyError: 'foo'

Moreover you can use list and Exception together in side_effect assignment. So the simplest way to do what you need is something like this:

>>> m = Mock(side_effect=[1, KeyError("bar"), 3])
>>> m("a")
1
>>> m("b")
 ...
KeyError: 'bar'
>>> m("c")
3

An other way to do it can be use a callable to write the logic of how your mock should react. In this case you are free to chose if it based on arguments or your test's state.

Answer 2

with a generator and Mock side_effect

from unittest.mock import Mock

def sample_generator():
    yield 1
    yield 2
    raise Exception()

gen = sample_generator()

def sideeffect():
    global gen
    for x in gen:
        return x


m = Mock(side_effect=sideeffect)
m() #1
m() #2
m() #3

you can get:

File "test.py", line 22, in <module>
   m() #3
   ....
   raise Exception()
Exception

I'm sure you can make it slightly cleaner with some effort, but it should solve your basic problem