How to get python's MagicMock to return a value the first time it is called and raise an exception the second time it is called?

Question

I have the following python program that uses mocking.

#!/usr/bin/env python 
import mock

def my_func1():
    return "Hello"

my_func = mock.MagicMock()
my_func.return_value = "Goodbye"

print my_func()
print my_func()

Output:

Goodbye
Goodbye

All is working as it should. Great.

But I want the mocked out method to return Goodbye the first time it is called and raise an exception the second time it is called. How can I do that??

score 7 · Accepted Answer · answered Sep 27 '17 at 05:24

As Sraw pointed out, you can use side_effect. I would probably use a generator function rather than introduce a global:

import mock

def effect(*args, **kwargs):
    yield "Goodbye"
    while True:
        yield Exception

my_func = mock.MagicMock()
my_func.side_effect = effect()

my_func() #Returns "Goodbye!'
my_func() #Raises exception
my_func() #Raises exception

Obviously you probably don't want to raise a bare Exception, but I wasn't sure what exception you were wanting to raise...

score 1 · Answer 2 · answered Sep 27 '17 at 05:14

You can use side_effect instead of return_value, for example:

import mock

a = 0
def my_func1():
    global a
    a += 1
    if a < 2:
        return "Goodbye"
    else:
        raise Exception()

my_func = mock.MagicMock()
my_func.side_effect = my_func1
print my_func()
# output Goodbye
print my_func()
# An exception was raised.

How to get python's MagicMock to return a value the first time it is called and raise an exception the second time it is called?

2 Answers2