Testing Django Management Commands with PyTest

, Jochen

The other day I wanted to test some Django management commands for django-cast and found this excellent post by Adam Johnson:

How to Unit Test a Django Management Command

But I use pytest to write my Django tests, so how would this be different?

A Management Command to Back up Media Files

Ok, so I wrote this little command that creates a backup of all media files. Django 4.2 introduced a new STORAGES setting, so it's now possible to write a backup command that works without having to know exactly how the media files are stored. It only assumes there are entries for production and backup storage backends and they are configured correctly. So it doesn't matter if your files are stored on S3 and the backup should be written to the local filesystem, or vice versa. The media_backup command does not need to know any of this.

from django.core.management.base import BaseCommand

from ...utils import storage_walk_paths
from .storage_backend import get_production_and_backup_storage


class Command(BaseCommand):
    help = (
        "backup media files from production to backup storage "
        "(requires Django >= 4.2 and production and backup storage configured)"
    )

    @staticmethod
    def backup_media_files(production_storage, backup_storage):
        for num, path in enumerate(storage_walk_paths(production_storage)):
            if not backup_storage.exists(path):
                with production_storage.open(path, "rb") as in_f:
                    backup_storage.save(path, in_f)
            if num % 100 == 0:  # pragma: no cover
                print(".", end="", flush=True)

    def handle(self, *args, **options):
        self.backup_media_files(*get_production_and_backup_storage())

Read on: TIL: Testing Django Management Commands with PyTest

Perhaps we should look at the get_production_and_backup_storage function first:

from django.core.management.base import BaseCommand

from ...utils import storage_walk_paths
from .storage_backend import get_production_and_backup_storage


class Command(BaseCommand):
    help = (
        "backup media files from production to backup storage "
        "(requires Django >= 4.2 and production and backup storage configured)"
    )

    @staticmethod
    def backup_media_files(production_storage, backup_storage):
        for num, path in enumerate(storage_walk_paths(production_storage)):
            if not backup_storage.exists(path):
                with production_storage.open(path, "rb") as in_f:
                    backup_storage.save(path, in_f)
            if num % 100 == 0:  # pragma: no cover
                print(".", end="", flush=True)

    def handle(self, *args, **options):
        self.backup_media_files(*get_production_and_backup_storage())

This function returns the production and backup media storage backends or raises an exception in the following two cases:

  • The installed version of Django is < 4.2 and therefore the storages dict does not exist yet
  • One of the required storage backends (production or backup) is missing

Let's have a look at the tests that make sure this happens:

import pytest

from django.core.management import call_command, CommandError



def test_media_backup_without_storages(settings):
    settings.STORAGES = {}
    with pytest.raises(CommandError) as err:
        call_command("media_backup")
    assert str(err.value) == "production or backup storage not configured"


def test_media_backup_with_wrong_django_version(mocker):
    mocker.patch("cast.management.commands.storage_backend.DJANGO_VERSION_VALID", False)
    with pytest.raises(CommandError) as err:
        call_command("media_backup")
    assert str(err.value) == "Django version >= 4.2 is required"

There's a test for each case. I could have printed out the error messages in the management command and used the capsys fixture to check that the output was right, but I thought raising a CommandError was a bit nicer here.

A more Complicated Test

Ok, the tests above are pretty easy to write. How about a test that mocks the storage backends to test the backup logic itself?

from collections.abc import Iterator
from contextlib import contextmanager
from io import BytesIO

import pytest
from django.core.files.storage import storages

from cast.management.commands.media_backup import Command as MediaBackupCommand


class StubStorage:
    def __init__(self) -> None:
        self._files: dict[str, BytesIO] = {}
        self._added: set[str] = set()

    def exists(self, path: str) -> bool:
        return path in self._files

    def was_added_by_backup(self, name: str) -> bool:
        return name in self._added

    def was_not_added_by_backup(self, name: str) -> bool:
        return name not in self._added

    def save(self, name: str, content: BytesIO) -> None:
        self.save_without_adding(name, content)
        self._added.add(name)

    def save_without_adding(self, name: str, content: BytesIO) -> None:
        self._files[name] = content

    def listdir(self, _path: str) -> tuple[list, dict[str, BytesIO]]:
        return [], self._files

    @contextmanager
    def open(self, name: str, _mode: str) -> Iterator[BytesIO]:
        try:
            yield self._files[name]
        finally:
            pass


@pytest.fixture
def stub_storages(settings):
    storage_stub = {"BACKEND": "tests.management_command_test.StubStorage"}
    settings.STORAGES = {"production": storage_stub, "backup": storage_stub}
    return storages


def test_media_backup_new_file_in_production(stub_storages):
    production, backup = stub_storages["production"], stub_storages["backup"]

    # given there's a new file added to production
    production.save_without_adding("foobar.jpg", BytesIO(b"foobar"))

    # when we run the backup command
    MediaBackupCommand().backup_media_files(production, backup)

    # then the file should have been added by the backup command
    assert backup.was_added_by_backup("foobar.jpg")

Return to blog