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())