9

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?

Nathan Tregillus
  • 6,006
  • 3
  • 52
  • 91

2 Answers2

29

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.

Michele d'Amico
  • 22,111
  • 8
  • 69
  • 76
  • seriously. I hate when people do that! I like your answer! this KeyError, is this a special method? I want this to throw a specific exception, is there a way to do that? – Nathan Tregillus Apr 16 '15 at 21:33
  • @NathanTregillus I used `KeyError` just as example. You can use every kind of `Exception` you want. – Michele d'Amico Apr 16 '15 at 21:35
  • 1
    @maazza yes of course... `m` is a mocked object. Even if you mean a _patched_ object the answer is yes again: `patch` patch an object with a `Mock` instance. – Michele d'Amico Mar 20 '19 at 08:01
  • When patching you need to define the side effects this way `m.side_effect = [1, KeyError(), 3]` – SimTae Feb 16 '21 at 18:38
4

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

user3012759
  • 1,977
  • 19
  • 22
  • shoot, this works to allow the first call, and throw on the second, but what if I want to allow a third call to be successful? – Nathan Tregillus Apr 15 '15 at 19:19
  • @NathanTregillus take a look to my answer it cover your comment too. – Michele d'Amico Apr 16 '15 at 06:53
  • I didn't edit you answer to cover the OP request just because I should change it too much. I did it in the past http://stackoverflow.com/a/28924186/4101725 . So filled a new answer. – Michele d'Amico Apr 16 '15 at 13:05
  • @Micheled'Amico feel free to edit in your points, it's simpler anyway and does cover both cases required by the OP. btw - down-vote is not work of mine... – user3012759 Apr 16 '15 at 13:16