Как использовать 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, а затем сопоставьте возвращаемые значения соответственно. Это идет вразрез с принципами СУХОЙ, и я хотел бы избежать этого.

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)