Python mocking elasticsearch connection generator `with` statement

user_6209161 0
Viewed 1060

I would like to test a class with a function that has with statement in it:

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

So I mocked the elstic_generator, and I mocked it get function when creating the class tested:

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()

But for some reason that doesn't work when using the with statement. However, if getting the connection without using with, like this:

x = self.__elastic_generator.get()

Then it works fine, and I get x = 'mocked elastic connection'.

So I guess the problem is related to some more functions calls being made when using the with and that I don't mock these functions.

Can someone please explain what happens under the hood and what else should I mock in order to be able to test it with the with statement?

Thanks.

python-mock python-2.7 with-statement mocking python
1 Answers

The with statement is a concept known as a context manager. A context manager has an __enter__ function for when you're entering the the with and an __exit__ function for when you're exiting the with (either by raising or by execution finishing inside the block).

The __enter__ function should return the value of whatever you're expecting to be assigned to the variable after as, which in this case would be es. So, to mock that, you don't want to mock the return value of .get(), you want to mock the return value of .get().__enter__(). That should look like this:

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