如何让Python中的MagicMock在第一次调用时返回一个值，在第二次调用时引发异常？

Question

我有一个使用Mocking的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，第二次被调用时引发异常。 我该怎么做？

- Saqib Ali

Answer 1

正如Sraw所指出的，您可以使用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，但我不确定您想引发哪个异常...

- Andrew Guy

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.

- Sraw