How to user Python Mock side effect to return different value based on function being mocked

Question

I have a common function that I want to use as a side effect for two different functions. Based on the function being mocked , they should return a different value. The code is:

def sleep_three(*args, **kwargs):
    logging.info('Sleeping 3 seconds')
    time.sleep(3)

Used as a side effect for two different functions:

 self.foo.side_effect = sleep_three
 self.bar.side_effect = sleep_three

For foo, I want a specific return value, whereas for bar I want a different return value. I know how to obtain different return values based on different arguments being passed in to side_effect function. In this case though both foo and bar are being passed in the same argument, they just need to return different values.

My options are:

1. Write two different functions for side_effect : sleep_three_foo and sleep_three_bar and then match up return values accordingly. This goes against DRY principles though and I would like to avoid it.

2. My question is : Is there a way within sleep_three to obtain the original function being mocked?

Something like

if original_mocked_function == foo:
    return a
elif original_mocked_function == bar:
    return b

Any pointers appreciated!

Answer 1

Define a function that takes the desired return value of the side-effect function as an argument and returns the appropriate function.

def make_sleeper(rv):
    def _(*args, **kwargs):
        logging.info("Sleeping 3 seconds")
        time.sleep(3)
        return rv
    return _

self.foo.side_effect = make_sleeper(5)
self.bar.side_effect = make_sleeper(9)