Customize test generation

To get started with generating tests, read Generate tests

Customizing Kolo’s test generation works by programmatically modifying the test plan with step or plan hooks.

For every test that Kolo generates from a trace, it first generates a test plan. And then b on that test plan the test code is generated. The test plan is a list of steps in JSON format (see the left side of the screenshot above)

View a list of all available steps: Test generation steps

Test generation step hooks

Kolo’s test generation works by generating a series of “steps”. Each of these steps represents a bit of the rendered test. Kolo provides the step_hook decorator which allows a step to be customised. For example, you could define a step hook that replaces a ModelCreate step with a FactoryCreate step:

from kolo import step_hook
from kolo.generate_tests.steps import ModelCreate, FactoryCreate


@step_hook(ModelCreate)
def model_create_to_factory(model_create):
    if model_create.module == "django.contrib.auth.models" and model_create.model == "User":
        factory_create = FactoryCreate(
            module="tests.factories",
            factory="UserFactory",
            fields=model_create.fields,
            defines_variable_name=model_create.defines_variable_name,
        )
        return [factory_create]
    return [model_create]

In the above code block, we received the original step (model_create) in our hook. Kolo then expects us to return a list of steps we’d like to replace that step with. In the case of our user model we’d like to replace it with just a single step, so we do return [factory_create].

---

Some step types come in two parts representing an indented block of code. For example, the TestFunction step is paired with the EndTestFunction step. A step hook for TestFunction will be passed a list of steps starting with a TestFunction, followed by all the steps for that test and finally an EndTestFunction step:

from kolo import step_hook
from kolo.generate_tests.steps import Code, EndTestFunction, TestFunction


@step_hook(TestFunction)
def add_pytest_fixture(steps):
    assert isinstance(steps[0], TestFunction)
    assert isinstance(steps[-1], EndTestFunction)

    steps[0].fixtures += ("my_cool_fixture",)
    steps.insert(1, Code("my_cool_fixture.do_something()"))
    return steps

Test generation plan hooks

It is also possible to define a hook that can work on the entire test plan. A plan hook takes a Plan instance and should return a new or modified Plan:

from kolo import plan_hook
from kolo.generate_tests.steps import EndTimeTravel, Code, StartTimeTravel, TestFunction


@plan_hook
def time_machine_decorator(plan):
    """
    Replace a time travel context manager with a decorator.

    First we find the `StartTimeTravel` step. If there are zero or
    multiple time travel steps, exit early without modifying the plan.

    Next we build up a new list of steps. This new list omits the
    `StartTimeTravel` and `EndTimeTravel` steps. It also inserts a new
    `Code` step defining a time travel decorator just before the
    `TestFunction` step.

    Finally the new steps are attached to the plan and the plan is
    returned.
    """
    time_travel_step = None
    for step in plan.steps:
        if isinstance(step, StartTimeTravel):
            if time_travel_step is None:
                time_travel_step = step
            else:
                # If there are two time travel context managers leave them both alone
                return plan

    if time_travel_step is None:
        # If there are no time travel context managers we have nothing to do
        return plan

    new_steps = []
    for step in plan.steps:
        if isinstance(step, (StartTimeTravel, EndTimeTravel)):
            pass
        elif isinstance(step, TestFunction):
            args = ", ".join(time_travel_step.args)
            time_travel_decorator = Code(
                f"@{time_travel_step.call}({args})",
                time_travel_step.imports,
            )
            new_steps.append(time_travel_decorator)
            new_steps.append(step)
        else:
            new_steps.append(step)

    plan.steps = new_steps
    return plan

Ensuring hooks are imported

Whether you’re using step hooks or plan hooks, you can ensure these are imported by Kolo by setting test_generation.hook_imports in .kolo/config.toml:

# .kolo/config.toml

[test_generation]

hook_imports = ["path.to.hooks.module", "path.to.more.hooks"]

Field parsers

Kolo supports many model fields provided by Django. Sometimes you may have a third-party field which isn’t being parsed correctly. In this case you can define a custom field parser:

def parse_custom_field(value, field):
    if field != "dotted.field.path.CustomField":
        return value

    # custom parsing logic

    return parsed_value

# .kolo/config.toml

[test_generation]

field_parsers = [
    "path.to.parse_custom_field",
]