How do I mock an open used in a with statement (using the Mock framework in Python)?

Question

How do I test the following code with mocks (using mocks, the patch decorator and sentinels provided by Michael Foord's Mock framework):

def testme(filepath):
    with open(filepath, 'r') as f:
        return f.read()

Answer 1

When used in a with statement, open returns a context manager. So the mock of open needs to also. A context manager must have an __enter__() and __exit__() method. The __enter__() method needs to return the file handle mock that will register the call to the read() method in the code above:

@patch('__builtin__.open')
def test_testme(self, open_mock):
    #
    # setup
    #
    context_manager_mock = Mock()
    open_mock.return_value = context_manager_mock
    file_mock = Mock()
    file_mock.read.return_value = sentinel.file_contents
    context_manager_mock.__enter__.return_value = file_mock

    #
    # exercise
    #
    result = cbot.testme(sentinel.filepath)

    #
    # verify
    #
    self.assertEquals(result, sentinel.file_contents)
    self.assertEquals(open_mock.call_args,
                      ((sentinel.filepath, 'r'), {}))
    self.assertEquals(context_manager_mock.method_calls,
                      [('__enter__', (), {}),
                       ('__exit__', (None, None, None), {})])
    self.assertEquals(file_mock.method_calls, [('read', (), {})])

Answer 2

Considering Mock() doesn't have __enter__ attribute set by default, how can the line

context_manager_mock.__enter__.return_value = file_mock

work?

How does with find __exit__ if it has not been declared yet?

Answer 3

Updated Daryl's answer to fix changes to Mock class.

@patch('__builtin__.open')
def test_testme(self, open_mock):
    #
    # setup
    #
    context_manager_mock = Mock()
    open_mock.return_value = context_manager_mock
    file_mock = Mock()
    file_mock.read.return_value = sentinel.file_contents
    enter_mock = Mock()
    enter_mock.return_value = file_mock
    exit_mock  = Mock()
    setattr( context_manager, '__enter__', enter_mock )
    setattr( context_manager, '__exit__', exit_mock )

    #
    # exercise
    #
    result = cbot.testme(sentinel.filepath)

    #
    # verify
    #
    self.assertEquals(result, sentinel.file_contents)
    self.assertEquals(open_mock.call_args,
                      ((sentinel.filepath, 'r'), {}))
    self.assertEquals(context_manager_mock.method_calls,
                      [('__enter__', (), {}),
                       ('__exit__', (None, None, None), {})])
    self.assertEquals(file_mock.method_calls, [('read', (), {})])