Python mock how to edit the returned value of a mocked function

Question

I have a function in a separate file and unittest in a different file, my goal is to get the actual return value of the function and edit it.

my_module.py

def function_to_test(a,b,c)(arg1,arg2,arg3):
    data_to_send = Mongoclient.find({_id:'arg1'})
    return data_to_send

def another_function():
    """Do something."""
    value_to_be_used = function_to_test(a,b,c)
    another_function_call_in_another_module(value_to_be_used)

test_file.py

class Mytest(unittest.TestCase):
    def test_one(self):
        # return_value is based on the original return value, 
        # and should vary based on the original returned value
        with patch(my_module.function_to_test, return_value='new data'):  
            my_module.another_function()

The new data is the original data + some edits to the original data.

How do I achieve this?

Answer 1

• You can keep a reference to the original function.
• And then you can define your patched function that calls the original one and adds some data.
• And then you can use your patched function as a side_effect for your mock.

class Mytest(unittest.TestCase):
    def test_one(self):
        # the return_value is based on the original return value and
        # should vary based on the original returned value

        original = my_module.function_to_test

        def patched(arg1, arg2, arg3):
            original_result = original(arg1, arg2, arg3)
            return original_result + ' @ new data'

        with patch("my_module.function_to_test", side_effect=patched):
            result = my_module.another_function()

        assert result == 'old data @ new data', result

def function_to_test(arg1, arg2, arg3):
    return 'old data'


def another_function_call_in_another_module(value_to_be_used):
    return value_to_be_used


def another_function():
    # Do something
    a, b, c = 1, 2, 3
    value_to_be_used = function_to_test(a, b, c)
    return another_function_call_in_another_module(value_to_be_used)

Some references:

• side_effect: https://docs.python.org/3/library/unittest.mock.html#unittest.mock.Mock.side_effect
• patch: https://docs.python.org/3/library/unittest.mock.html#unittest.mock.patch