模拟上下文管理器中使用的类

Question

I have a module method that uses a context manager with a Class that offers iteration:

import module

def ReadLog(log_file):
  with module.LogReader(log_file) as reader:
    combined_list = []
    for record in reader:
      combined_list.append(record)
    return combined_list

In my unittest, I'm trying to mock module.LogReader such that it produces records defined by myself. This is where I'm at so far:

import mock
import my_mod

@mock.patch.object(my_mod.module, 'LogReader')
def testReadLog(self, mock_logreader):
  filename = '/fake/file/name.log'
  my_mod.ReadLog(filename)

  # Verify that module.LogReader was called
  mock_logreader.assert_called_once_with(filename)

But I haven't been able to make the iteration produce records so far. The LogReader class has __enter__, __exit__, __iter__, next methods, and I've tried doing mock_logreader.return_value = '123' for example but this results in the error message AttributeError: __exit__ .

What sauce am I missing here?

Answer 1

You need to set return_value on the functions you want to mock, before you call the function you're testing. Note that on anything where you haven't set this, it'll return a new MagicMock . You can intercept this by setting a return_value on this new MagicMock .

So for example:

file = 'foo' 
mock = MagicMock()  # could be gotten through patch
mock(file).__enter__().__iter__.return_value = [1,2,3,4,5]

with mock(file) as reader:
    for x in reader:
        print(x)

See the Python data model for how the magic functions work.