如何根据模拟的函数使用Python Mock副作用返回不同的值

时间:2019-03-25 19:23:49

标签: python unit-testing python-mock

我有一个常用功能,希望用作两个不同功能的副作用。基于被模拟的函数,它们应该返回不同的值。 代码是:

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

任何指针表示赞赏!

1 个答案:

答案 0 :(得分: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)