Search code examples
pythonunit-testingmockingpython-unittest

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?


Solution

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