How do I test a function called inside the tarfile.open context manager?

Question

I've written a function that creates a tarball.

# tarball.py
import os
import tarfile


def create_tarball():
    with tarfile.open("path/to/tar/file", "w:gz") as tar:
        tar.add(
            "/path/to/included/directory",
            arcname=os.path.basename("/path/to/included/directory"),
        )

I've written a passing test with an assertion on the tarfile.open context manager.

from unittest.mock import patch

from tarball import create_tarball


@patch("tarball.tarfile.open")
def test_create_tarball_partial(mock_open):
    create_tarball()

    mock_open.assert_called_with("path/to/tar/file", "w:gz")

How do I write a test for a function called inside the tarfile.open context manager?

Ryan Payne · Accepted Answer · 2019-12-06T02:42:35.917

To write a test for the tar.add function, do the following.

Patch the os.path.basename function
Create a mock of tar.add using MagicMock
Set the value of the tar.add function inside the context manager to the MagicMock¹
Set the return value of the os.path.basename function
Assert on the mocked os.path.basename and tarfile.open functions

For example:

from unittest.mock import patch, MagicMock

from tarball import create_tarball


@patch("tarball.os.path.basename")
@patch("tarball.tarfile.open")
def test_create_tarball_full(mock_open, mock_basename):
    mock_add = MagicMock()
    mock_open.return_value.__enter__.return_value.add = mock_add
    mock_basename.return_value = "/path/to/included/directory"

    create_tarball()

    mock_open.assert_called_with("path/to/tar/file", "w:gz")
    mock_basename.assert_called_with("/path/to/included/directory")
    mock_add.assert_called_with(
        "/path/to/included/directory", arcname="/path/to/included/directory"
    )

¹Python: Mocking a context manager

How do I test a function called inside the tarfile.open context manager?

1 Answers1