Как заставить Python MagicMock возвращать значение при первом вызове и вызывать исключение при втором вызове?

Question

У меня есть следующая программа Python, которая использует насмешки.

#!/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()

Выход:

Goodbye
Goodbye

Все работает как надо. Отлично.

Но я хочу, чтобы макетный метод вернулся Goodbye в первый раз он вызывается и вызывает исключение во второй раз, когда он вызывается. Как мне это сделать??

Answer 1

Как указал Сроу, вы можете использовать side_effect, Я бы, вероятно, использовал функцию генератора, а не вводил глобальную:

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

Очевидно, вы, вероятно, не хотите поднять голый Exception, но я не был уверен, какое исключение вы хотели поднять...

Answer 2

Ты можешь использовать side_effect вместо return_value, например:

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.