Create a gzip file like object for unit testing

Viewed 1053

I want to test a Python function that reads a gzip file and extracts something from the file (using pytest).

import gzip

def my_function(file_path):
    output = []

    with gzip.open(file_path, 'rt') as f:
        for line in f:
            output.append('something from line')

    return output

Can I create a gzip file like object that I can pass to my_function? The object should have defined content and should work with gzip.open()

I know that I can create a temporary gzip file in a fixture but this depends on the filesystem and other properties of the environment. Creating a file-like object from code would be more portable.

3 Answers

You can use the io and gzip libraries to create in-memory file objects. Example:

import io, gzip

def inmem():
    stream = io.BytesIO()
    with gzip.open(stream, 'wb') as f:
        f.write(b'spam\neggs\n')
    stream.seek(0)
    return stream

You should never try to test outside code in a unit test. Only test the code you wrote. If you're testing gzip, then gzip is doing something wrong (they should be writing their own unit tests). Instead, do something like this:

from unittest import mock

@mock.Mock('gzip', return_value=b'<whatever you expect to be returned from gzip>')
def test_my_function(mock_gzip):
    file_path = 'testpath'
    output = my_function(file_path=file_path)
    mock_gzip.open.assert_called_with(file_path)
    assert output == b'<whatever you expect to be returned from your method>'

That's your whole unit test. All you want to know is that gzip.open() was called (and you assume it works or else gzip is failing and that's their problem) and that you got back what you expected from the method being tested. You specify what gzip returns based on what you expect it to return, but you don't actually call the function in your test.

It's a bit verbose but I'd do something like this (I have assumed that you saved my_function to a file called patch_one.py):

import patch_one  # this is the file with my_function in it
from unittest.mock import patch
from unittest import TestCase


class MyTestCase(TestCase):
    def test_my_function(self):
        # because you used "with open(...) as f", we need a mock context
        class MyContext:
            def __enter__(self, *args, **kwargs):
                return [1, 2]  # note the two items

            def __exit__(self, *args, **kwargs):
                return None

        # in case we want to know the arguments to open()
        open_args = None

        def f(*args, **kwargs):
            def my_open(*args, **kwargs):
                nonlocal open_args
                open_args = args
                return MyContext()

            return my_open

        # patch the gzip.open in our file under test
        with patch('patch_one.gzip.open', new_callable=f):  

            # finally, we can call the function we want to test
            ret_val = patch_one.my_function('not a real file path')

            # note the two items, corresponding to the list in __enter__()
            self.assertListEqual(['something from line', 'something from line'], ret_val)

            # check the arguments, just for fun
            self.assertEqual('rt', open_args[1])

If you want to try anything more complicated, I would recommend reading the unittest mock docs because how you import the "patch_one" file matters as does the string you pass to patch().

There will definitely be a way to do this with Mock or MagicMock but I find them a bit hard to debug so I went the long way round.

Related