How to make a python mocked out function return a specific value conditional on an argument to the function?

Question

I have a python 2.7x Tornado application that when run serves up a handful of RESTful api endpoints.

My project folder includes numerous test cases that rely on the python mock module such as shown below.

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'})
        )

The RESTful endpoint http://localhost/my_service/my_endpoint has two internal calls to my_method respectively: my_method(my_arg=1) and my_method(my_arg=2).

I want to mock out my_method in this test-case such that it returns 0 if it is called with my_arg==2, but otherwise it should return what it would always normally return. How can I do it?

I know that I should do something like this:

mock_my_method.return_value = SOMETHING

But I don't know how to properly specify that something so that its behavior is conditional on the arguments that my_method is called with. Can someone show me or point me to an example??

Answer 1

I want to mock out my_method in this test-case such that it returns 0 if it is called with my_arg==2, but otherwise it should return what it would always normally return. How can I do it?

Write your own method mock calling the original one on condition:

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)

For patching: if you don't need to assert how often the mocked method was called and with what args etc, it is sufficient to pass the mock via new argument:

@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)

If you want to invoke the whole mock assertion machinery, use side_effect as suggested in the other answer. Example:

@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)

Answer 2

you could use side_effect to change return value dynamically:

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()

---

update: as you want to run original my_method code under some condition, you may need a method proxy, Mock can't get back the real function object being patched.

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)

will outputs:

my_arg 1 ret 10000
my_arg 2 ret 0