# How to: Customize test generation

To get started with generating tests, read [How to: Generate tests](generate-tests)

<img width="1798" alt="image" src="https://github.com/user-attachments/assets/63b6b5e4-d1e9-4d45-97aa-f35e1b4af237">


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](../reference/steps.md)

### 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:

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

    # custom parsing logic

    return parsed_value
```

```toml
# .kolo/config.toml

[test_generation]

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