Refactoring Incident Status Transitions with Pure Functions and Unit Tests

Hello! Last lesson, you used branch coverage to find untested decisions and then wrote assertions for the behavior those branches should produce. Now you’ll put those tests to work during a refactor: separating an incident status decision from the code that saves the incident and sends a notification.

This is a practical way to improve Python code that began life as a script. You can keep its observable behavior while making the business rule small enough to test directly. The example below is self-contained; adapt its names and statuses if your incident project already has equivalent models.

Find the decision inside the workflow

Suppose your incident project has three statuses. An incident may change from NEW to IN_PROGRESS, then from IN_PROGRESS to RESOLVED. All other requests, including requests to set the current status again, are rejected. This is the rule for this example, not a universal incident-management policy.

Start with these domain types in src/incident_app/incident_domain.py:

from dataclasses import dataclass
from enum import Enum


class Status(str, Enum):
    NEW = "new"
    IN_PROGRESS = "in_progress"
    RESOLVED = "resolved"


@dataclass
class Incident:
    id: int
    status: Status


class InvalidStatusTransition(ValueError):
    pass

Imagine that src/incident_app/incident_workflow.py currently contains both the rule and its consequences:

from collections.abc import Callable

from incident_app.incident_domain import (
    Incident,
    InvalidStatusTransition,
    Status,
)


def change_status(
    incident: Incident,
    target: Status,
    *,
    save: Callable[[Incident], None],
    notify: Callable[[int, Status], None],
) -> None:
    if incident.status == Status.NEW:
        allowed = target == Status.IN_PROGRESS
    elif incident.status == Status.IN_PROGRESS:
        allowed = target == Status.RESOLVED
    else:
        allowed = False

    if not allowed:
        raise InvalidStatusTransition(
            f"Cannot change status from {incident.status.value} "
            f"to {target.value}"
        )

    incident.status = target
    save(incident)
    notify(incident.id, incident.status)

The decision depends only on the current status and the requested status. Saving and notifying are different responsibilities: they affect the outside world and need collaborators. When all three are in one function, a test of the decision must also supply those collaborators.

Watch the short extraction example in How To Write Unit Tests For Existing Python Code // Part 2 of 2 by ArjanCodes. The example extracts a calculation from an object; focus on why the calculation becomes simpler to test when it no longer needs that object.

There is a broader design reason for doing this. In Testing State Machines, the author uses an application with hardware actions, but the same testing difficulty occurs when a status rule is mixed with database writes or messages. Read only the conceptual opening; you do not need the article’s later implementation example.

Protect the existing workflow first

Before moving code, record what a caller can observe. In particular, a successful change should update the incident, save it, and notify with its new status. A rejected change should leave the incident untouched and call neither collaborator.

Add tests/test_incident_workflow.py:

import pytest

from incident_app.incident_domain import (
    Incident,
    InvalidStatusTransition,
    Status,
)
from incident_app.incident_workflow import change_status


@pytest.mark.parametrize(
    ("current", "target"),
    [
        (Status.NEW, Status.IN_PROGRESS),
        (Status.IN_PROGRESS, Status.RESOLVED),
    ],
)
def test_allowed_change_updates_and_reports_status(
    current: Status, target: Status
) -> None:
    incident = Incident(id=42, status=current)
    events: list[tuple[str, int, Status]] = []

    def save(saved: Incident) -> None:
        events.append(("save", saved.id, saved.status))

    def notify(incident_id: int, status: Status) -> None:
        events.append(("notify", incident_id, status))

    change_status(incident, target, save=save, notify=notify)

    assert incident.status == target
    assert events == [
        ("save", 42, target),
        ("notify", 42, target),
    ]


def test_rejected_change_has_no_effects() -> None:
    incident = Incident(id=42, status=Status.NEW)

    def unexpected_call(*args: object) -> None:
        pytest.fail("A rejected change must not save or notify")

    with pytest.raises(
        InvalidStatusTransition, match="new.*resolved"
    ):
        change_status(
            incident,
            Status.RESOLVED,
            save=unexpected_call,
            notify=unexpected_call,
        )

    assert incident.status == Status.NEW

The events list represents observable calls at the workflow boundary. It checks their order because this example’s existing workflow saves before notifying; it does not assert which if or elif ran. Run the tests now, before changing the implementation:

python -m pytest -q tests/test_incident_workflow.py

Passing baseline tests give you a reference point. They do not yet cover every status pair; the focused decision tests below will specify the complete small rule.

Extract and test the pure decision

A pure function makes its result depend only on its arguments and produces no external effects. Here, the proposed contract is:

next_status(current: Status, requested: Status) -> Status

For an allowed request it returns the requested status. For a rejected request it raises InvalidStatusTransition with details a caller can use. It does not mutate an Incident, save anything, notify anyone, or read hidden configuration or time. Given the same inputs, its result—or its rejection—is predictable.

Write its tests in tests/test_incident_transitions.py. The two allowed pairs and seven rejected pairs account for every combination of the three statuses under our stated policy:

import pytest

from incident_app.incident_domain import (
    InvalidStatusTransition,
    Status,
)
from incident_app.incident_transitions import next_status


@pytest.mark.parametrize(
    ("current", "requested"),
    [
        (Status.NEW, Status.IN_PROGRESS),
        (Status.IN_PROGRESS, Status.RESOLVED),
    ],
)
def test_allowed_transition_returns_new_status(
    current: Status, requested: Status
) -> None:
    assert next_status(current, requested) == requested


@pytest.mark.parametrize(
    ("current", "requested"),
    [
        (Status.NEW, Status.NEW),
        (Status.NEW, Status.RESOLVED),
        (Status.IN_PROGRESS, Status.NEW),
        (Status.IN_PROGRESS, Status.IN_PROGRESS),
        (Status.RESOLVED, Status.NEW),
        (Status.RESOLVED, Status.IN_PROGRESS),
        (Status.RESOLVED, Status.RESOLVED),
    ],
)
def test_rejected_transition_identifies_both_statuses(
    current: Status, requested: Status
) -> None:
    with pytest.raises(
        InvalidStatusTransition,
        match=f"{current.value}.*{requested.value}",
    ):
        next_status(current, requested)

Create src/incident_app/incident_transitions.py. You could begin with a stub that returns requested and run the tests: the allowed cases pass, while the rejected cases expose its missing rule. Then implement the function:

from incident_app.incident_domain import (
    InvalidStatusTransition,
    Status,
)


def next_status(current: Status, requested: Status) -> Status:
    if current == Status.NEW:
        allowed = requested == Status.IN_PROGRESS
    elif current == Status.IN_PROGRESS:
        allowed = requested == Status.RESOLVED
    else:
        allowed = False

    if not allowed:
        raise InvalidStatusTransition(
            f"Cannot change status from {current.value} "
            f"to {requested.value}"
        )

    return requested

Notice the package boundaries: both the workflow and the decision function import shared domain types, while the decision function does not import the workflow. That keeps the extracted rule independent and avoids a circular import.

The tests assert results and meaningful rejection details, rather than reproducing the conditionals in their assertions. For further guidance on that distinction, read the following passage from My Python Testing Style Guide by Thea Flowers.

Tests alone cannot prove purity: a function could return the right value while secretly logging or modifying global state. Its interface, implementation, and tests work together here. The interface needs only two enum values; inspection confirms that the implementation has no mutable input or external collaborator.

Reconnect the rule without moving the effects

Replace change_status() in incident_workflow.py with a version that delegates only the decision. Keep its existing signature:

from collections.abc import Callable

from incident_app.incident_domain import Incident, Status
from incident_app.incident_transitions import next_status


def change_status(
    incident: Incident,
    target: Status,
    *,
    save: Callable[[Incident], None],
    notify: Callable[[int, Status], None],
) -> None:
    new_status = next_status(incident.status, target)

    incident.status = new_status
    save(incident)
    notify(incident.id, new_status)

The order matters. next_status() runs before mutation or either external call, so a rejected request still has no effects. For a valid request, the workflow remains responsible for updating the incident, saving it, and notifying. The pure function is not a replacement for the workflow; it is the workflow’s testable decision-making part.

Run the whole test suite after the change:

python -m pytest -q

If you use the branch-coverage setup from the previous lesson, you can rerun it against the refactored package. Use its report to investigate any missing decisions, but keep judging the tests by the behavior they assert. The small transition function is now quick to test without a database, environment setup, or test doubles.

This refactor preserves the workflow’s current behavior; it does not solve every persistence concern. For example, if save() fails after the in-memory incident has changed, restoring that object or managing a database transaction requires a separate design decision. Keeping that question distinct is part of maintaining a clear boundary.

Takeaways

You protected observable behavior first, extracted the status rule into next_status(), tested all status pairs for this small policy, and left mutation and external calls in change_status(). The result is not fewer tests at all costs: it is tests placed at the right boundaries—direct, exhaustive tests for the decision and a smaller set of workflow tests for its effects.

This closes the pytest module. Next, you’ll begin shaping the incident-management API by mapping resources and operations to REST routes, HTTP methods, and status codes.