简体   繁体   中英

Mock file stream in python

I want to test my logging module. I have a function that generates a line and write it to file. I have the following code:

def open_file_stream(filename, opt):
    return open(filename, opt)

def close_file_stream(f):
    f.close()

def write_line_to_log_file(message, pid, module, state, data={}, datatype='AKV'):
    file_path = '/path/to/file'
    try:
        f = open_file_stream(file_path, 'a')
        # generate_line only takes the arguments and formats it to string
        f.write(generate_line(message, pid, module, state, data, 'AKV') + '\n')
    except Exception as e:
        raise RuntimeError('Failed to write log line to %s: %s' % (file_path, str(e)))
    finally:
        f.close()

I tried testing function write_line_to_log_file with following code:

class Unit_test(unittest.TestCase):
    @patch('my_logger.generate_line', return_value='abc')
    def test_write_line_to_log_file(self, arg1):
        with patch('my_logger.open_file_stream', return_value=StringIO.StringIO()) as f:
            my_logger.write_line_to_log_file(message='message123', pid=1234, 
                                             module='abc/def', state='OK', 
                                             data={'test_key': 'test_value'})
            should_be = 'abc'
            f.seek(0)
            self.assertEqual(f.readline(), should_be)

I thought that patch will create real object that I could read and check if it writes the content to filestream as expected, instead I got:

AssertionError: <MagicMock name='open_file_stream.readline()' id='3066213676'> != 'abc'

How should I fix that? Or if the content is in this MagicMock object how should I read it?

First of all you should mock out external resources not your own code. In this case open . You can test open using mock_open from the mock library format.

Let's say you have this code here in your my_module.py

def write_contents(contents):
    try:
        with open('/my/path', 'w') as fh:
            fh.write(contents)
    except IOError:
        print "Something wrong happened"

You can have the following test case using mock_open

import mock
import unittest

import my_module

class TestClass(unittest.TestCase):
    def test_my_module(self):
        m = mock.mock_open()
        with mock.patch('my_module.open', m, create=True) as mocked_open: 
            contents = "Sample write content"
            my_module.write_contents(contents)
            file_handle = mocked_open()
            # Do the assert
            file_handle.write.assert_called_once_with(contents, 'w')

This way you are telling python that during the running of the test, assume a new open method exists and assert that it behaves as expected.

When you do with patch('my_logger.open_file_stream', return_value=StringIO.StringIO()) as f: , f is a mock object that will only contain things you add to it ie by assigning the return_value attribute something or giving it a side_effect.

You can tune the above example to your case. The file_handle object will behave similar to any other open object, with the methods such as write , read , readinto etc and you can test as you wish.

The technical post webpages of this site follow the CC BY-SA 4.0 protocol. If you need to reprint, please indicate the site URL or the original address.Any question please contact:yoyou2525@163.com.

 
粤ICP备18138465号  © 2020-2024 STACKOOM.COM