Legacy Lobotomy — Creating Management Commands for Seeding Assignment Categories and Targets

This content originally appeared on Level Up Coding – Medium and was authored by Yevhen Nosolenko

Legacy Lobotomy — Creating Management Commands for Seeding Assignment Categories and Targets

This is the 14th tutorial in our series on refactoring a legacy Django project. In the previous part, we created management commands to generate teams and users. In this tutorial, we will continue building up our test data by adding commands for seeding assignment categories and targets.

Categories help organize different types of assignments, and targets define which group of users an assignment should go to. This structure is important for how assignments are created and distributed in the system. By generating this data automatically, we make it easier to test and work with the project during development.

To get more ideas regarding the project used for this tutorial or check other tutorials in this series, check the introductory article published earlier.

Legacy Lobotomy — Confident Refactoring of a Django Project

Let’s get started by creating Django management commands that will add some categories and targets to our database.

Origin branch and destination branch

1. If you want to follow the steps described in this tutorial, you should start from the seeding-teams-and-users branch.
2. If you want to only see final code state after all changes from this tutorial are applied, you can check the seeding-assignment-categories-and-targets branch.

Preparing the assignments app

In this section, we will be working with the assignments app. But first, we need to make a couple of changes to its structure.

Step 1. Create a few new packages within the assignments app: factories, management and tests.

Step 2. Add a new commands package to the assignments/management folder.

Step 3. Add a new management package to the assignments/tests folder.

Step 4. Add a new commands package to the assignments/tests/management folder.

The final structure of the assignments app should be as shown in the image below:

Creating a command for seeding categories

We already have commands to generate teams and users. Next, we are going to create commands for seeding assignments. An important part of preparing assignments is creating categories.

Command requirements

The requirements for our category seeding command are as follows:

• The command should generate categories with unique names.
• It should accept an optional argument –count to specify the number of categories to be created. If no argument is provided, it should create 5 categories by default.
• It must validate the –count argument, ensuring it's a positive integer.
• The command should output the IDs of the newly created categories to standard output for future reference.

Let’s start by creating categories, which will be used in our future assignment seeding commands.

Step 1: Create and register the category factory

Create a new category.py module inside the assignments/factories directory with the following content:

import uuid

import factory

from assignments.models import Category


class CategoryFactory(factory.django.DjangoModelFactory):
    class Meta:
        model = Category

    name = factory.LazyFunction(lambda: f'Fake Category #{uuid.uuid4().hex}')

Next, make sure this factory is imported in the assignments/factories/__init__.py module:

from .category import CategoryFactory

Finally, register the factory in src/conftest.py so it can be used in tests:

from assignments.factories import CategoryFactory

# Register assignment factories
register(CategoryFactory)

This makes two fixtures available in tests:

• category — provides a pre-built Category instance that can be injected into test functions
• category_factory — gives you full control to generate categories using factory methods like .create() or .build()

Step 2: Define the category seeding command

Create a new seed_categories.py module with the following content in the assignments/management/commands directory:

from django.core.management import BaseCommand, CommandError

from assignments.factories import CategoryFactory

class Command(BaseCommand):
    help = 'Seed fake categories.'

    def add_arguments(self, parser):
        parser.add_argument(
            '--count',
            required=False,
            type=int,
            default=5,
            help='The number of categories which should be created.'
        )

    def handle(self, count, *args, **options):
        if not isinstance(count, int) or count < 1:
            raise CommandError('The --count argument must be an integer value greater than or equal to 1.')

        categories = CategoryFactory.create_batch(count)
        category_ids = ','.join(str(category.pk) for category in categories)
        self.stdout.write(category_ids)

Step 3: Implement tests

Add a new test module test_seed_categories.py with the content as follows to the assignments/tests/management/commands folder:

import pytest
from django.core.management import CommandError, call_command

from assignments.models import Category

@pytest.mark.django_db
class TestSeedCategories:
    def test_when_count_argument_is_not_provided_then_5_categories_should_be_created(self):
        call_command('seed_categories')

        assert Category.objects.count() == 5

    @pytest.mark.parametrize('count', [-10, 0, 0.5, 'invalid'])
    def test_when_count_argument_is_invalid_then_error_should_be_thrown(self, count):
        with pytest.raises(CommandError) as exc_info:
            call_command('seed_categories', count=count)

        assert str(exc_info.value) == 'The --count argument must be an integer value greater than or equal to 1.'

    @pytest.mark.parametrize('count', [1, 5, 12])
    def test_when_count_argument_is_provided_then_requested_number_of_categories_should_be_created(self, count):
        call_command('seed_categories', count=count)

        assert Category.objects.count() == count

    @pytest.mark.parametrize('count', [1, 5, 12])
    def test_when_category_is_created_then_its_id_should_be_written_to_standard_output(self, count, capsys):
        call_command('seed_categories', count=count)

        captured_output = capsys.readouterr().out.strip()
        captured_category_ids = {int(category_id) for category_id in captured_output.split(',')}

        assert len(captured_category_ids) == count

Step 4: Run and verify

Run the tests and ensure they all pass successfully.

Next, run the following command from the root of the project and make sure that 5 categories are successfully created:

python src/manage.py seed_categories

Now, run the following command from the root of the project and verify that 8 additional categories are successfully created:

python src/manage.py seed_categories --count 8

The image below shows the output from running both commands:

That’s all for seeding categories. Let’s move on to creating a command to seed assignment targets.

Creating a command for seeding assignment targets

Before we can generate assignments themselves, we need some AssingmentTarget (note the typo in the model name — this is how the model was originally named) records in place. These targets define the demographic slices (age ranges, gender flags, etc.) to which each assignment will apply.

Command requirements

Our seeding command will:

• Generate unique targets with random min_age/max_age values within valid bounds and a random combination of gender flags (male, female, non_binary, transgender, other).
• Accept an optional argument –count to specify the number of categories to be created. If no argument is provided, the command must create 5 assignment targets by default.
• Validate that –count is a positive integer — otherwise it raises a CommandError.
• Print the newly created target IDs to standard output so they can be captured for later use.

Let’s implement this command.

Step 1: Create and register the assignment target factory

Create a module assignment_target.py with the following content in the assignments/factories folder:

import uuid

import factory

from assignments.models import AssingmentTarget


class AssignmentTargetFactory(factory.django.DjangoModelFactory):
    class Meta:
        model = AssingmentTarget

    name = factory.LazyFunction(lambda: f'Fake Target #{uuid.uuid4().hex}')
    min_age = factory.Faker('pyint', min_value=13, max_value=99)
    max_age = factory.Faker('pyint', min_value=factory.SelfAttribute('..min_age'), max_value=99)
    male = factory.Faker('pybool')
    female = factory.Faker('pybool')
    non_binary = factory.Faker('pybool')
    transgender = factory.Faker('pybool')
    other = factory.Faker('pybool')
    law_explorer = factory.Faker('pybool')

Here, we used factory.SelfAttribute(‘..min_age’) to reference the factory’s min_age field, ensuring that max_age is always greater than or equal to min_age.

Next, import the factory in assignments/factories/__init__.py:

from .assignment_target import AssignmentTargetFactory

After that register the factory in src/conftest.py to be able to use it in tests:

from .factories import AssignmentTargetFactory

register(AssignmentTargetFactory, 'assignment_target')

We used a custom name, assignment_target, because the model name, AssingmentTarget, contains a typo. Without this, the objects injected into tests would inherit the typo and use a name like assingment_target.

Step 2: Define the assignment target seeding command

Create a new module seed_assignment_targets.py with the following content in the assignments/management/commands directory:

from django.core.management import BaseCommand, CommandError
from assignments.factories import AssignmentTargetFactory

class Command(BaseCommand):
    help = 'Seed fake assignment targets.'

    def add_arguments(self, parser):
        parser.add_argument(
            '--count',
            type=int,
            default=5,
            help='The number of assignment targets which should be created.'
        )

    def handle(self, count, *args, **options):
        if not isinstance(count, int) or count < 1:
            raise CommandError('The --count argument must be an integer value greater than or equal to 1.')

        assignment_targets = AssignmentTargetFactory.create_batch(count)
        ids = ','.join(str(target.pk) for target in assignment_targets)
        self.stdout.write(ids)

Step 3: Implement tests

Place test_seed_assignment_targets.py with the following content in the assignments/tests/management/commands directory:

import pytest
from django.core.management import CommandError, call_command

from assignments.models import AssingmentTarget


@pytest.mark.django_db
class TestSeedAssignmentTargets:
    def test_when_count_argument_is_not_provided_then_5_assignment_targets_should_be_created(self):
        call_command('seed_assignment_targets')

        assert AssingmentTarget.objects.count() == 5

    @pytest.mark.parametrize('count', [-10, 0, 0.5, 'invalid'])
    def test_when_count_argument_is_invalid_then_error_should_be_thrown(self, count):
        with pytest.raises(CommandError) as exc_info:
            call_command('seed_assignment_targets', count=count)

        assert str(exc_info.value) == 'The --count argument must be an integer value greater than or equal to 1.'

    @pytest.mark.parametrize('count', [1, 5, 12])
    def test_when_count_argument_is_provided_then_requested_number_of_assignment_targets_should_be_created(self, count):
        call_command('seed_assignment_targets', count=count)

        assert AssingmentTarget.objects.count() == count

    @pytest.mark.parametrize('count', [1, 5, 12])
    def test_when_assignment_target_is_created_then_its_id_should_be_written_to_standard_output(self, count, capsys):
        call_command('seed_assignment_targets', count=count)

        captured_output = capsys.readouterr().out.strip()
        captured_assignment_target_ids = {int(target_id) for target_id in captured_output.split(',')}

        assert len(captured_assignment_target_ids) == count

Step 4: Run and verify

Run the tests and ensure they all pass successfully.

Next, run the following command from the root of the project and confirm that 5 targets are created:

python src/manage.py seed_assignment_targets

Then, run the command again with the –count argument to verify that 8 additional targets are added:

python src/manage.py seed_assignment_targets - count 8

The image below shows the output from running the command with and without –count:

The creation of the seed_assignment_targets command can be considered complete.

Conclusion

In this tutorial, we created Django management commands to seed assignment categories and targets. We built new factories where needed and followed a structure that keeps the codebase clean and extendable. With this data in place, we’re ready to move on to the next step. In the following tutorial, we’ll generate assignments and connect them to categories and targets.

Legacy Lobotomy — Creating Management Commands for Seeding Assignment Categories and Targets was originally published in Level Up Coding on Medium, where people are continuing the conversation by highlighting and responding to this story.

This content originally appeared on Level Up Coding – Medium and was authored by Yevhen Nosolenko