如何使python模拟出的函数返回以函数的参数为​​条件的特定值？

Question

我有一个python 2.7x Tornado应用程序，该应用程序在运行时可提供一些RESTful api端点。

我的项目文件夹包括许多依赖python mock模块的测试用例，如下所示。

from tornado.testing import AsyncHTTPTestCase
from mock import Mock, patch
import json
from my_project import my_model

class APITestCases(AsyncHTTPTestCase):

    def setUp(self):
        pass

    def tearDown(self):
        pass

    @patch('my_project.my_model.my_method')
    def test_something(
        self,
        mock_my_method
    ):

        response = self.fetch(
            path='http://localhost/my_service/my_endpoint',
            method='POST',
            headers={'Content-Type': 'application/json'},
            body=json.dumps({'hello':'world'})
        )

RESTful端点http://localhost/my_service/my_endpoint分别有两个内部对my_method的调用：my_method(my_arg=1)和my_method(my_arg=2)。

在这个测试用例中，我想模拟my_method，以便如果用0 == 2调用它会返回my_arg，但是否则它应该返回它将始终返回的值正常返回。我该怎么办？

我知道我应该做这样的事情：

mock_my_method.return_value = SOMETHING

但是我不知道如何正确地指定某种东西，以使其行为取决于调用my_method的参数。有人可以给我展示或给我指出一个例子吗？

Answer 1

您可以使用side_effect动态更改返回值：

class C:
    def foo(self):
        pass  

def drive():
    o = C()
    print(o.foo(my_arg=1))
    print(o.foo(my_arg=2))  

def mocked_foo(*args, **kwargs):
    if kwargs.get('my_arg') == 2:
        return 0
    else:
        return 1

@patch('__main__.C.foo')
def test(mock):
    mock.side_effect = mocked_foo
    drive()

---

更新：由于您想在某些条件下运行原始my_method代码，因此可能需要方法代理，Mock无法取回被修补的真实函数对象。

from unittest.mock import patch

class MyClass:

    def my_method(self, my_arg):
        return 10000

def func_wrapper(func):
    def wrapped(*args, **kwargs):
        my_arg = kwargs.get('my_arg')
        if my_arg == 2:
            return 0
        return func(*args, **kwargs)
    return wrapped

def drive(o, my_arg):
    print('my_arg', my_arg, 'ret', o.my_method(my_arg=my_arg))

def test():
    with patch.object(MyClass, 'my_method', new=func_wrapper(MyClass.my_method)):
        o = MyClass()
        drive(o, 1)
        drive(o, 2)

将输出：

my_arg 1 ret 10000
my_arg 2 ret 0

Answer 2

  

在这个测试用例中，我想模拟my_method，使得如果用my_arg==2调用它会返回0，但是否则它应该返回通常返回的值。我该怎么办？

编写您自己的方法模拟，条件是调用原始方法：

from my_project import my_model

my_method_orig = my_project.my_model.my_method
def my_method_mocked(self, *args, my_arg=1, **kwargs):
    if my_arg == 2:  # fake call
        return 0
    # otherwise, dispatch to real method
    return my_method_orig(self, *args, **kwargs, my_arg=my_arg)

用于修补：如果您不需要断言模拟方法的调用频率以及args等参数，则只需通过new参数传递模拟即可：

@patch('my_project.my_model.my_method', new=my_method_mocked)
def test_something(
    self,
    mock_my_method
):

    response = self.fetch(...)
    # this will not work here:
    mock_my_method.assert_called_with(2)

如果要调用整个模拟断言机制，请按照其他答案中的建议使用side_effect。示例：

@patch('my_project.my_model.my_method', side_effect=my_method_mocked, autospec=True)
def test_something(
    self,
    mock_my_method
):

    response = self.fetch(...)
    # mock is assertable here 
    mock_my_method.assert_called_with(2)