如何根據模擬的函數使用Python Mock副作用返回不同的值

Question

我有一個常用功能，希望用作兩個不同功能的副作用。 基於被模擬的函數，它們應該返回不同的值。 代碼是：

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

用作兩種不同功能的副作用：

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

對於foo，我想要一個特定的返回值，而對於bar，我想要一個不同的返回值。 我知道如何基於傳遞給side_effect函數的不同參數來獲取不同的返回值。 在這種情況下，盡管foo和bar都在同一個參數中傳遞，但它們只需要返回不同的值。

我的選擇是：

1. 為side_effect編寫兩個不同的函數：sleep_three_foo和sleep_three_bar，然后相應地匹配返回值。 但是，這違反了DRY原則，我想避免這樣做。

2. 我的問題是：sleep_three中是否有辦法獲得被嘲笑的原始函數？

就像是

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

任何指針表示贊賞！

Answer 1

定義一個函數，該函數將副作用函數的期望返回值作為參數並返回適當的函數。

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)