0

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!

peaxol
  • 497
  • 4
  • 13
  • How different *are* `sleep_three_foo` and `sleep_three_bar`? You could write one function that generates the correct function based on some argument, then write `self.foo.side_effect = make_mock(...); self.bar.side_effect = make_mock(...)`. – chepner Mar 25 '19 at 19:32

1 Answers1

3

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)
chepner
  • 497,756
  • 71
  • 530
  • 681