Testing

Axio ships with testing helpers in axio.testing that make it easy to write fast, deterministic tests for agents, tools, and custom components.

StubTransport

StubTransport is a fake transport that yields pre-configured event sequences instead of calling a real LLM:

from axio.testing import StubTransport, make_text_response

transport = StubTransport(
    [
        make_text_response("Hello!"),
    ]
)
assert len(transport._responses) == 1
assert transport._call_count == 0

Each entry in the list is one transport call. The stub cycles through them in order, repeating the last one if the agent makes more calls than expected.

A sequence is any list[StreamEvent], not only what the factories below build. A refusal, a reasoning signature or a ProviderEvent can therefore be fed to an agent with no network and no API key.

What StubTransport does not do: stream() ignores messages, tools and system entirely. A test that needs to assert on what the agent sent — that a stored ReasoningBlock was replayed, that a tool was withheld — has to record it itself. See “Asserting what the transport was sent” below.

Factory functions

make_text_response

Create an event sequence for a simple text reply. iteration defaults to 2 here and to 1 in make_tool_use_response, so the pair reads as one tool turn followed by the answer:

from axio import TextDelta, IterationEnd, StopReason, Usage
from axio.testing import make_text_response

events = make_text_response(text="Hello world", iteration=1)
assert events == [
    TextDelta(index=0, delta="Hello world"),
    IterationEnd(
        iteration=1,
        stop_reason=StopReason.end_turn,
        usage=Usage(input_tokens=10, output_tokens=5),
    ),
]

make_tool_use_response

Create an event sequence for a tool call:

from axio import ToolUseStart, ToolInputDelta, IterationEnd, StopReason
from axio.testing import make_tool_use_response

events = make_tool_use_response(
    tool_name="greet",
    tool_id="call_1",
    tool_input={"name": "Alice"},
    iteration=1,
)
assert len(events) == 3
assert isinstance(events[0], ToolUseStart)
assert events[0].name == "greet"
assert events[0].tool_use_id == "call_1"
assert isinstance(events[1], ToolInputDelta)
assert "Alice" in events[1].partial_json
assert isinstance(events[2], IterationEnd)
assert events[2].stop_reason == StopReason.tool_use

make_stub_transport

Create transport that returns a single “Hello world” text response:

from axio import TextDelta, IterationEnd
from axio.testing import make_stub_transport

transport = make_stub_transport()
assert len(transport._responses) == 1
assert isinstance(transport._responses[0][0], TextDelta)
assert isinstance(transport._responses[0][-1], IterationEnd)

make_ephemeral_context

Create a fresh in-memory context store:

from axio import MemoryContextStore
from axio.testing import make_ephemeral_context

context = make_ephemeral_context()
assert isinstance(context, MemoryContextStore)
assert context.session_id is not None

make_echo_tool

Create a test tool that echoes its input as JSON:

from axio.testing import make_echo_tool

tool = make_echo_tool()
assert tool.name == "echo"
assert "JSON" in tool.description

Testing an agent with tools

A typical test sets up a stub that first requests a tool call, then returns text after seeing the result:

import asyncio
from axio import Agent
from axio.testing import (
    StubTransport,
    make_tool_use_response,
    make_text_response,
    make_ephemeral_context,
    make_echo_tool,
)


async def test_agent_calls_tool():
    transport = StubTransport(
        [
            make_tool_use_response("echo", tool_input={"text": "hello"}),
            make_text_response("Done!"),
        ]
    )
    agent = Agent(
        system="You are a test agent.",
        tools=[make_echo_tool()],
        transport=transport,
    )
    result = await agent.run("Say hello", make_ephemeral_context())
    assert result == "Done!"


asyncio.run(test_agent_calls_tool())

No @pytest.mark.asyncio decorator needed - the project uses asyncio_mode = "auto".

Asserting what the transport was sent

StubTransport throws away its arguments, so record them with a transport of your own. This is how a transport author tests reasoning replay, and how anyone tests that the selector withheld a tool:

import asyncio
from collections.abc import AsyncIterator
from typing import Any
from axio import (
    Agent,
    MemoryContextStore,
    StopReason,
    StreamEvent,
    TextDelta,
    Tool,
    Usage,
)
from axio.blocks import TextBlock
from axio.events import IterationEnd
from axio.messages import Message


class Recorder:
    """A transport that keeps every history and tool list it was given."""

    def __init__(self, responses: list[list[StreamEvent]]) -> None:
        self.responses = responses
        self.sent: list[list[Message]] = []
        self.offered: list[list[Tool[Any]]] = []

    def stream(
        self, messages: list[Message], tools: list[Tool[Any]], system: str
    ) -> AsyncIterator[StreamEvent]:
        self.sent.append(list(messages))
        self.offered.append(list(tools))
        return self._replay(
            self.responses[min(len(self.sent) - 1, len(self.responses) - 1)]
        )

    async def _replay(self, events: list[StreamEvent]) -> AsyncIterator[StreamEvent]:
        for event in events:
            yield event


async def main() -> None:
    transport = Recorder(
        [
            [
                TextDelta(index=0, delta="hi"),
                IterationEnd(
                    iteration=0, stop_reason=StopReason.end_turn, usage=Usage(10, 2)
                ),
            ]
        ]
    )
    agent = Agent(system="", transport=transport)
    await agent.run("hello", MemoryContextStore())

    [message] = transport.sent[0]
    said = "".join(b.text for b in message.content if isinstance(b, TextBlock))
    # The user turn is already in the history the transport is handed, and the agent prefixes
    # it with a local timestamp.
    assert said.endswith("hello")


asyncio.run(main())

Testing a declined turn

A refusal is a normal terminal outcome, not an exception. Build one directly and assert on both the text and the stop reason:

import asyncio
from axio import Agent, MemoryContextStore, Refusal, StopReason, Usage
from axio.events import IterationEnd, SessionEndEvent
from axio.testing import StubTransport


async def main() -> None:
    transport = StubTransport(
        [
            [
                Refusal(index=0, text="I can't help with that.", category="policy"),
                IterationEnd(
                    iteration=0, stop_reason=StopReason.refusal, usage=Usage(12, 5)
                ),
            ]
        ]
    )
    agent = Agent(system="", transport=transport)

    # get_final_text() collects refusal text, so this is no longer the empty string.
    assert await agent.run("...", MemoryContextStore()) == "I can't help with that."

    ends = [
        e
        async for e in agent.run_stream("...", MemoryContextStore())
        if isinstance(e, SessionEndEvent)
    ]
    assert ends[0].stop_reason is StopReason.refusal


asyncio.run(main())

Testing an SSE reader

A Reader claims only the events it interprets, and forwards the rest. The test worth writing is that every name it claims is one the provider publishes, plus a strict read that fails the day the provider sends something new:

import pytest
from collections.abc import Iterator
from dataclasses import dataclass
from axio.events import StreamEvent, TextDelta
from axio_sse import EVENT_NAME, Decoder, Reader, UnknownEvent, Wire, on

# In a real test this is the provider's published event list, checked into the repository.
PUBLISHED = frozenset({"text.delta", "ping", "web_search.started"})


@dataclass(frozen=True, slots=True)
class TextChunk(Wire, name="text.delta"):
    text: str = ""


class Reply(Reader[StreamEvent], by=EVENT_NAME):
    @on(TextChunk)
    def _text(self, wire: TextChunk) -> Iterator[StreamEvent]:
        yield TextDelta(index=0, delta=wire.text)


# One direction only. The reader is meant to be a subset: it names what it interprets.
assert Reply.names() <= PUBLISHED

[event] = Decoder().decode(b"event: something.new\ndata: {}\n\n", final=True)
with pytest.raises(UnknownEvent):
    Reply().read(event, strict=True)

Testing tools in isolation

Test a tool handler directly:

import asyncio
from axio import Tool


async def word_count(text: str) -> str:
    """Count words in text."""
    count = len(text.split())
    return f"The text contains {count} words."


async def test_word_count():
    tool = Tool(name="word_count", handler=word_count)
    result = await tool(text="one two three")
    assert "3" in result


asyncio.run(test_word_count())

Or test through the Tool wrapper to exercise guards:

import asyncio
from axio import Tool


async def word_count(text: str) -> str:
    """Count words in text."""
    count = len(text.split())
    return f"The text contains {count} words."


async def test_word_count_tool():
    tool = Tool(name="word_count", handler=word_count)
    result = await tool(text="one two three")
    assert "3" in result


asyncio.run(test_word_count_tool())

Testing guards

import asyncio
import pytest
from typing import Any
from axio import Tool, PermissionGuard, GuardError


async def word_count(text: str) -> str:
    """Count words."""
    return str(len(text.split()))


_tool: Tool[Any] = Tool(name="word_count", handler=word_count)


class MaxLengthGuard(PermissionGuard):
    def __init__(self, max_length: int = 10000) -> None:
        self.max_length = max_length

    async def check(self, tool: Tool[Any], **kwargs: Any) -> dict[str, Any]:
        for name, value in kwargs.items():
            if isinstance(value, str) and len(value) > self.max_length:
                raise GuardError(f"Field '{name}' exceeds {self.max_length} characters")
        return kwargs


async def test_guard_allows_short_input():
    guard = MaxLengthGuard(max_length=100)
    result = await guard.check(_tool, text="short")
    assert result == {"text": "short"}


async def test_guard_denies_long_input():
    guard = MaxLengthGuard(max_length=5)
    with pytest.raises(GuardError):
        await guard.check(_tool, text="this is way too long")


asyncio.run(test_guard_allows_short_input())
asyncio.run(test_guard_denies_long_input())

Testing context stores

import asyncio
from axio import MemoryContextStore
from axio.messages import Message
from axio.blocks import TextBlock


async def test_context_append_and_history():
    ctx = MemoryContextStore()
    msg = Message(role="user", content=[TextBlock(text="hello")])
    await ctx.append(msg)
    history = await ctx.get_history()
    assert len(history) == 1
    assert history[0].role == "user"


asyncio.run(test_context_append_and_history())