CODEWITHSUNDEEP
pythonunit-testingtestingmocking
Benjamin Hodgson
Benjamin Hodgson

Reputation: 44634

Python's 'mock' - responding to an argument

I need to test a method that opens two files and writes different data to each of them. It doesn't matter what order the files get written in.

Here's how I'd test a method that only needs to open one file, using a Mock to replace open:

from io import BytesIO
import mock

class MemorisingBytesIO(BytesIO):
    """Like a BytesIO, but it remembers what its value was when it was closed."""
    def close(self):
        self.final_value = self.getvalue()
        super(MemorisingBytesIO, self).close()

open_mock = mock.Mock()
open_mock.return_value = MemorisingBytesIO()

with mock.patch('__builtin__.open', open_mock):
    write_to_the_file()  # the function under test

open_mock.assert_called_once_with('the/file.name', 'wb')
assert open_mock.return_value.final_value == b'the data'

I'm having trouble modifying this approach to work with a method that writes to two files. I've considered using side_effect to return two MemorisingBytesIOs sequentially, and asserting that each of them contains the right data, but then the test will be brittle: if the order of the calls in the method changes, the test will fail.

So what I really want to do is to have open_mock return one MemorisingBytesIO when it's called with one file name, and a different one when it's called with the other. I've seen this in other languages' mocking libraries: is it possible in Python without subclassing Mock?

Upvotes: 3

Views: 792

Answers (2)

falsetru
falsetru

Reputation: 369094

How about following approach? (Use class attribute to hold file content):

from io import BytesIO
import mock

class MemorisingBytesIO(BytesIO):
    """Like a BytesIO, but it remembers what its value was when it was closed."""
    contents = {}
    def __init__(self, filepath, *args, **kwargs):
        self.filepath = filepath
        super(MemorisingBytesIO, self).__init__()
    def close(self):
        self.contents[self.filepath] = self.getvalue()
        super(MemorisingBytesIO, self).close()

def write_to_the_file():
    with open('a/b.txt', 'wb') as f:
        f.write('the data')
    with open('a/c.txt', 'wb') as f:
        f.write('another data')


#MemorisingBytesIO.contents.clear()
open_mock = mock.Mock(side_effect=MemorisingBytesIO)
with mock.patch('__builtin__.open', open_mock):
    write_to_the_file()  # the function under test

open_mock.assert_called_once_with('a/b.txt', 'wb')
open_mock.assert_called_once_with('a/c.txt', 'wb')
assert MemorisingBytesIO.contents['a/b.txt'] == b'the data'
assert MemorisingBytesIO.contents['a/c.txt'] == b'another data'

Upvotes: 1

Benjamin Hodgson
Benjamin Hodgson

Reputation: 44634

I have since discovered the way to do what I originally wanted using mock. You can set side_effect equal to a function; when the mock is called, that function is passed the arguments.

In [1]: import mock

In [2]: def print_it(a, b):
   ...:     print b
   ...:     print a
   ...:     

In [3]: m = mock.Mock(side_effect=print_it)

In [4]: m('hello', 2)
2
hello

So here's the how you'd write the original example to work with two files:

fake_file_1 = MemorisingBytesIO()
fake_file_2 = MemorisingBytesIO()

def make_fake_file(filename, mode):
    if filename == 'a/b.txt':
        return fake_file_1
    elif filename == 'a/c.txt':
        return fake_file_2
    else:
        raise IOError('Wrong file name, Einstein')

open_mock = mock.Mock(side_effect=make_fake_file)
with mock.patch('__builtin__.open', open_mock):
    write_to_the_file()

assert ('a/b.txt', 'wb') in open_mock.call_args
assert ('a/c.txt', 'wb') in open_mock.call_args
assert fake_file_1.final_value == 'file 1 data'
assert fake_file_2.final_value == 'file 2 data'

Upvotes: 0

Related Questions