How to: Customize test generation

To get started with generating tests, read How to: 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",
]