Mocking a class used in a with statement

Question

I have a class which has an __exit__ and __enter__ function so that I can use it in a with statement, e.g.:

with ClassName() as c:
    c.do_something()

I am now trying to write a unit test to test this. Basically, I am trying to test that do_something() has only been called once.

An example (which I called testmocking1):

class temp:
    def __init__(self):
        pass

    def __enter__(self):
        pass

    def __exit__(self, exc_type, exc_val, exc_tb):
        pass

    def test_method(self):
        return 1


def fun():
    with temp() as t:
        return t.test_method()

And my test:

import unittest
import test_mocking1
from test_mocking1 import fun
import mock
from mock import patch

class MyTestCase(unittest.TestCase):
    @patch('test_mocking1.temp', autospec = True)
    def test_fun_enter_called_once(self, mocked_object):
        fun()
        mocked_object.test_method.assert_called_once()

if __name__ == '__main__':
    unittest.main()

So I would expect this to pass, because the test_method has been called exactly once in the function fun(). But the actual result that I get is:

======================================================================
FAIL: test_fun_enter_called_once (__main__.MyTestCase)
----------------------------------------------------------------------
Traceback (most recent call last):
  File "<path_to_virtual_env>\lib\site-packages\mock\mock.py", line 1305, in patched
    return func(*args, **keywargs)
  File "<File_with_test>", line 11, in test_fun_enter_called_once
mocked_object.test_method.assert_called_once()
  File "<path_to_virtual_env>\lib\site- 
packages\mock\mock.py", line 915, in assert_called_once
    raise AssertionError(msg)
AssertionError: Expected 'test_method' to have been called once. Called 0 times.

How do I test whether a function in a class which is created using a with statement has been called (either once or multiple times), and (related) how do I set the results of those calls (using .side_effect or .return_value)?

Answer 1

The with statement takes whatever __enter__ returns to bind to the name in the as <name> part. You bound it to t:

with temp() as t:
    t.test_method()

Note that temp() is called, so the with statement starts with temp.return_value. t is not temp.return_value either, it is whatever temp().__enter__() returns, so you need to use the return value for that call:

entered = mocked_object.return_value.__enter__.return_value
entered.test_method.assert_called_once()

Extending on this, if you want to alter what test_method() returns, do so on the return value of mocked_object.return_value.__enter__.return_value.

You can always print out the mock_calls() attribute of your object to see what has happened to it:

>>> from test_mocking1 import fun
>>> from mock import patch
>>> with patch('test_mocking1.temp', autospec = True) as mocked_object:
...     fun()
...
>>> print(mocked_object.mock_calls)
[call(),
 call().__enter__(),
 call().__enter__().test_method(),
 call().__exit__(None, None, None)]
>>> mocked_object.return_value.__enter__.return_value.test_method.called
True
>>> mocked_object.return_value.__enter__.return_value.test_method.call_count
1

Note that your actual implementation of temp.__enter__() returns None, so without mocking your fun() function fails with an attribute error.