How to Unit Test a Django Management Command

Here comes the test train!

This post is an adapted extract from my book Speed Up Your Django Tests, available now.

When we write custom management commands, it’s easy to write integration tests for them with call_command(). This allows us to invoke the management command as it runs under manage.py, and retrieve the return code, standard output, and standard error. It’s great, but has some overhead, making our tests slower than necessary. If we have logic separated out of the command’s handle() method, it improves both readability and testability, as we can unit test it separately.

Example Command

Take this management command, that implements some book title normalization rules:

from django.core.management.base import BaseCommand
from django.db.transaction import atomic

from example.core.models import Book


class Command(BaseCommand):
    help = "Normalize all Book titles"

    def add_arguments(self, parser):
        parser.add_argument(
            "--write",
            action="store_true",
            default=False,
            help="Actually edit the database",
        )

    @atomic
    def handle(self, *args, write, **kwargs):
        if not write:
            self.stdout.write("In dry run mode (--write not passed)")

        books = Book.objects.select_for_update()
        for book in books:
            book.title = self.normalize_title(book.title)
        if write:
            Book.objects.bulk_update(books, ["title"])
            self.stdout.write(f"Updated {len(books)} book(s)")

    def normalize_title(self, title):
        if not title:
            return "Unknown"

        if title[0].islower():
            title = title[0].upper() + title[1:]

        if title.endswith("."):
            title = title[:-1]

        title = title.replace("&", "and")

        return title

(These is the same logic as the form in How to Unit Test a Django Form.)

The logic inside normalize_title() has been separated from the handle() method, allowing us to use and test it in isolation.

Integration Tests

We could write integration tests for this command that create data, run the command, and check the output was expected and the data was correctly updated:

from io import StringIO

from django.core.management import call_command
from django.test import TestCase

from example.core.models import Book


class NormalizeBookTitlesTests(TestCase):
    def call_command(self, *args, **kwargs):
        out = StringIO()
        call_command(
            "normalize_book_titles",
            *args,
            stdout=out,
            stderr=StringIO(),
            **kwargs,
        )
        return out.getvalue()

    def test_dry_run(self):
        book_empty = Book.objects.create(title="")

        out = self.call_command()
        self.assertEqual(out, "In dry run mode (--write not passed)\n")

        book_empty.refresh_from_db()
        self.assertEqual(book_empty.title, "")

    def test_write_empty(self):
        book_empty = Book.objects.create(title="")

        out = self.call_command("--write")

        self.assertEqual(out, "Updated 1 book(s)\n")
        book_empty.refresh_from_db()
        self.assertEqual(book_empty.title, "Unknown")

    def test_write_lowercase(self):
        book_lowercase = Book.objects.create(title="lowercase")

        out = self.call_command("--write")

        self.assertEqual(out, "Updated 1 book(s)\n")
        book_lowercase.refresh_from_db()
        self.assertEqual(book_lowercase.title, "Lowercase")

    def test_write_full_stop(self):
        book_full_stop = Book.objects.create(title="Full Stop.")

        out = self.call_command("--write")

        self.assertEqual(out, "Updated 1 book(s)\n")
        book_full_stop.refresh_from_db()
        self.assertEqual(book_full_stop.title, "Full Stop")

    def test_write_ampersand(self):
        book_ampersand = Book.objects.create(title="Dombey & Son")

        out = self.call_command("--write")

        self.assertEqual(out, "Updated 1 book(s)\n")
        book_ampersand.refresh_from_db()
        self.assertEqual(book_ampersand.title, "Dombey and Son")

These tests capture the command output through the stdout argument to call_command(), and then make assertions on it.

The visible repetition in writing and reading the Book instances to the database point us to some overhead. There’s also the overhead of running call_command() to test only behaviour from normalize_title(). Seeing both of these nudges us to move the normalization tests to unit tests.

Unit Tests

We can change some of these tests to directly test normalize_title() instead:

from io import StringIO

from django.core.management import call_command
from django.test import SimpleTestCase, TestCase

from example.core.management.commands.normalize_book_titles import Command
from example.core.models import Book


class NormalizeBookTitlesSimpleTests(SimpleTestCase):
    def test_normalize_empty(self):
        result = Command().normalize_title("")
        self.assertEqual(result, "Unknown")

    def test_normalize_lowercase(self):
        result = Command().normalize_title("lowercase")
        self.assertEqual(result, "Lowercase")

    def test_normalize_full_stop(self):
        result = Command().normalize_title("Full Stop.")
        self.assertEqual(result, "Full Stop")

    def test_normalize_ampersand(self):
        result = Command().normalize_title("Dombey & Son")
        self.assertEqual(result, "Dombey and Son")


class NormalizeBookTitlesTests(TestCase):
    def call_command(self, *args, **kwargs):
        call_command(
            "normalize_book_titles",
            *args,
            stdout=StringIO(),
            stderr=StringIO(),
            **kwargs,
        )

    def test_dry_run(self):
        book_empty = Book.objects.create(title="")

        self.call_command()

        book_empty.refresh_from_db()
        self.assertEqual(book_empty.title, "")

    def test_write(self):
        book_empty = Book.objects.create(title="")

        self.call_command("--write")

        book_empty.refresh_from_db()
        self.assertEqual(book_empty.title, "Unknown")

Because the new tests don’t touch the database, we have placed them in a SimpleTestCase class, saving some database overhead. They’re also shorter as they don’t require any arrangement step.

There are two integration tests left, to cover the “dry run” and “write” pathways.

Fin

May you write faster, more targeted tests,

—Adam


Read my book Boost Your Git DX to Git better.


Subscribe via RSS, Twitter, Mastodon, or email:

One summary email a week, no spam, I pinky promise.

Related posts:

Tags: ,