Python模拟elasticsearch连接生成器`with`语句

Question

我想测试一个包含with 语句的函数的类：

func_to_test():
   ....
   with self.__elastic_generator.get() as es:
      print 'about to use the es connection....'

所以我模拟了 elstic_generator，并在创建测试类时模拟了 get 函数：

elastic_gen = mock.Mock()
elstic_gen.get = mock.Mock()
elastic_gen.get.side_effect = ['mocked elastic connection']
tested_class = TestedClass(elastic_gen)
tested_class.func_to_test()

但由于某些原因，使用 with 语句时不起作用。 但是，如果在不使用with 的情况下获得连接，像这样：

x = self.__elastic_generator.get()

然后它工作正常，我得到x = 'mocked elastic connection'。

所以我猜这个问题与使用 with 时进行的更多函数调用有关，并且我没有模拟这些函数。

有人可以解释一下幕后发生了什么，我还应该模拟什么才能使用with 语句对其进行测试？

谢谢。

Answer 1

with 语句是一个称为上下文管理器的概念。上下文管理器有一个__enter__ 函数，用于当您进入with 时和一个__exit__ 函数，用于当您退出with 时（通过提升或在块内执行完成）。

__enter__ 函数应返回您期望分配给as 之后的变量的值，在本例中为es。所以，要模拟它，你不想模拟.get() 的返回值，你想模拟.get().__enter__() 的返回值。应该是这样的：

elastic_gen = mock.Mock()
elastic_gen.return_value.__enter__.return_value = 'mocked elastic connection'