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

Python mocking elasticsearch connection generator `with` statement

我想用一个包含 with 语句的函数测试 class:

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

所以我嘲笑了elstic_generator,我在创建class测试时嘲笑了它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 语句对其进行测试?

谢谢。

with 语句是一个称为上下文管理器的概念。上下文管理器在您进入 with 时有一个 __enter__ 函数,在您退出 with 时有一个 __exit__ 函数(通过提高或通过块内执行完成)。

__enter__ 函数应该 return 您希望在 as 之后分配给变量的任何值,在本例中为 es .因此,要模拟它,您不想模拟 .get() 的 return 值,而是想要模拟 .get().__enter__() 的 return 值。那应该是这样的:

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