Testing TramAI Code

What it is: tramai-testing provides deterministic in-memory providers, failure simulation, request recording, and a fluent assertion DSL, so you can test AI-dependent application code with no network calls. In 0.6.0 it also ships the shared contract TCKs — ProviderTck and the store TCKs — used to prove that providers and persistence stores conform to TramAI's behavioural contracts.

When to use it: Every test that involves a TramAI service. Never hit a live model in CI.

What The Module Provides

Application test doubles

Contract TCKs and test infrastructure

The TCKs live in the module's testFixtures source set. They are for extension and module authors, not for application tests.

Dependency

This is a published artifact, managed by the BOM. Declare the platform once, then the module without a version:

dependencies {
    testImplementation(platform("dev.tramai:tramai-bom:0.6.0"))
    testImplementation("dev.tramai:tramai-testing")
}

Basic Test Pattern

val provider = MockAiProvider {
    onMethod("analyze") respondWith """{"status":"ok"}"""
}
val observer = RecordingOperationObserver()

val tramai = Tramai {
    provider(provider, default = true)
    model("gpt-4o", "mock")
    observer(observer)
}

val service = tramai.create<Analyzer>()
val result = runBlocking { service.analyze("invoice-1") }

assertEquals(Status("ok"), result)

TramaiAssertions.assertThat(provider, observer)
    .whenCalled("analyze")
    .wasCalledTimes(1)
    .andRetried(0)
    .andParsedSuccessfully()
    .emittedProvider("mock")

MockAiProvider keys responses by the service method name (request.operationMethod), and consumes them in order on repeated calls. providerId() is always "mock".

Assertion API

Testing Structured Output

Mock responses feed straight into the structured-output pipeline — extraction, parse, shape validation, value validation, and repair retries.

data class Status(val status: String)

@AiService
interface Analyzer {
    @Operation(prompt = "Analyze the invoice", model = "gpt-4o")
    suspend fun analyze(invoiceId: String): Status
}

@Test
fun `parses structured output correctly`() = runBlocking {
    val provider = MockAiProvider {
        onMethod("analyze") respondWith """{"status":"approved"}"""
    }

    val tramai = Tramai {
        provider(provider, default = true)
        model("gpt-4o", "mock")
    }

    assertEquals(Status("approved"), tramai.create<Analyzer>().analyze("inv-1"))
}

Simulating Retries

Queue several responses for the same method; each call consumes the next one.

val provider = MockAiProvider {
    onMethod("analyze") respondWith "not json"          // attempt 1 → parse failure
    onMethod("analyze") respondWith """{"status":"ok"}""" // attempt 2 → success
}

Assert the retry explicitly:

TramaiAssertions.assertThat(provider, observer)
    .whenCalled("analyze")
    .wasCalledTimes(2)
    .andRetried(1)
    .andParsedSuccessfully()

When retries are exhausted, the call fails with StructuredOutputException, carrying a typed StructuredOutputFailureCode such as REPAIR_EXHAUSTED and an attemptCount. Raw prompt, raw response, and validation detail are deliberately not exposed on the public exception — classification is typed, not textual.

assertThatThrownBy { runBlocking { service.analyze("inv-1") } }
    .isInstanceOfSatisfying(StructuredOutputException::class.java) { error ->
        assertEquals(3, error.attemptCount)
        assertEquals(StructuredOutputFailureCode.REPAIR_EXHAUSTED, error.failureCode)
    }

Simulating Provider Failures

Use SimulatedFailureProvider for provider-level failures:

val provider = SimulatedFailureProvider {
    onMethod("summarize").retryableFailure("rate limited", statusCode = 429)
    onMethod("summarize") respondWith "recovered summary"
}

Non-retryable failures surface immediately with their typed classification:

val provider = SimulatedFailureProvider {
    onMethod("summarize").nonRetryableFailure("unauthorized", statusCode = 401)
}

// ...
assertThatThrownBy { runBlocking { service.summarize("inv-1") } }
    .isInstanceOfSatisfying(ProviderException::class.java) { error ->
        assertEquals(401, error.statusCode)
        assertFalse(error.retryable)
    }

This is the right way to verify retryable-failure handling, non-retryable propagation, and recovery after transient provider errors.

Recording And Inspecting Requests

RecordingOperationObserver captures per-attempt metadata. MockAiProvider also records the raw requests it received, because it implements RecordedRequestProvider:

val sent = provider.requests.single()
println(sent.model)            // requested model
println(sent.operationMethod)  // "analyze"
println(sent.messages)         // the messages the engine built

That is the shortest path to asserting prompt and argument wiring.

Tool-Calling Tests

MockTool builds typed tool doubles without a real implementation, and records every input it receives:

val tool: MockTool<CustomerQuery, Customer> = MockTool.fixed(
    name = "customer_lookup",
    description = "Look up a customer",
    response = Customer(id = "c-1", name = "Ada"),
)

Register it with tools(tool) on the builder. MockTool implements TramaiTool, so it is indistinguishable from a real tool to the engine; tool.calls holds the inputs that reached it, which makes argument-plumbing assertions straightforward.

Contract TCKs For Extension Authors

If you implement a ModelProvider or a persistence store, the shared TCK is the contract — not an optional extra. Enrolment is architecture-enforced, so a missing runner fails the build.

./gradlew :tramai-testing:test

Runs every *EnrollmentArchitectureTest gate, which requires each concrete implementation to have a valid, correctly named TCK runner.

Walk-throughs:

• Adding a Provider — implement ModelProvider and enrol in ProviderTck
• Adding a Store — implement a store SPI and enrol in its TCK
• Adding an Approval State — the approval TCKs and lifecycle oracles
• Custom Providers — the SPI from the inside

The TCKs are deterministic and offline by design: ProviderTck answers every request from canned StubHttpClient responses, and store TCKs drive implementations through harnesses with MutableClock / MutableMillisClock rather than wall-clock sleeps.

What To Test

Recommended TramAI-facing application tests:

• prompt and argument wiring through service methods
• structured return types, including repair and exhaustion paths
• retry and fallback behaviour
• provider routing decisions
• business logic wrapped around TramAI service calls

Most teams start with three tests:

1. the happy path returns the typed result you expect;
2. a malformed first response retries and then succeeds;
3. a retryable provider failure recovers, and a non-retryable one fails the way you expect.

What Not To Over-Test

• Every internal prompt word choice — test the contract, not the prose.
• Provider HTTP specifics in application tests; those belong in provider TCK tests.
• Jackson deserialisation itself.
• Model intelligence or answer quality.

Limitations

• MockAiProvider does not stream. It implements ModelProvider, not StreamCapable. Test streaming against a real local model, or use a provider adapter's stub seam.
• MockAiProvider does not synthesise tool calls. Use MockTool for tool-calling shapes, or a purpose-built provider double.
• No response generator for application tests. Property-based and model-based testing exists where it pays for itself — in the TCK lifecycle oracles and in the engine's retry/fallback lifecycle tests — not as a general-purpose application-testing feature.
• Mock providers are per-test. Complex multi-provider failover chains must be orchestrated explicitly.

Next Steps

• Observability — assert spans and metrics with an in-memory exporter
• CI Integration — the gate set TramAI itself runs
• Structured Output — the parse-and-repair contract
• Providers & Routing — model routing and fallbacks