Mock return value of a function for one case input, otherwise return original value

apps.my_module.py

def my_func(name: str) -> str:
    if name == 'a':
        return 'name is a'
    elif name == 'b':
        return 'name is b'
    else:
        return 'not defined'

I want to mock my_func in tests.py the mocked function return value for a is name is AA not name is a, but the return value for b and else must stay the same

I want to do this using unittest, how to do this?

Answer

The solution goes as the following:

in tests.py

from apps import my_module
import unitteest

class test(unittest.TestCase):
    def any_test(self, *_):
        original_func = my_module.my_func
        def mocked_func(name):
            if name == 'a':
                return 'name is AA'
            else:
                return original_func(name)