| 1 | // Package testutil provides reusable test helpers for the agent package. |
| 2 | // MockProvider replaces a real LLM backend in agent tests with scripted |
| 3 | // responses, request recording, and error injection — so the harness |
| 4 | // loop, cache behaviour, and tool dispatch can be verified without |
| 5 | // network calls. |
| 6 | package testutil |
| 7 | |
| 8 | import ( |
| 9 | "context" |
| 10 | "fmt" |
| 11 | "sync" |
| 12 | |
| 13 | "reasonix/internal/provider" |
| 14 | ) |
| 15 | |
| 16 | // Turn describes one expected Stream call: the text, optional reasoning, |
| 17 | // optional tool calls, usage telemetry, and optionally an error to inject. |
| 18 | type Turn struct { |
| 19 | Text string |
| 20 | Reasoning string |
| 21 | ToolCalls []provider.ToolCall |
| 22 | Usage *provider.Usage |
| 23 | // Chunks, when non-empty, is emitted exactly as provided. It is useful for |
| 24 | // edge cases such as partial tool-call starts followed by an error. |
| 25 | Chunks []provider.Chunk |
| 26 | |
| 27 | // StreamError, when set, causes Stream to return this error before any |
| 28 | // chunks, simulating a network or auth failure for that turn. |
| 29 | StreamError error |
| 30 | // ChunkError, when set, is emitted after the scripted chunks, simulating a |
| 31 | // mid-stream provider failure after partial output has reached the agent. |
| 32 | ChunkError error |
| 33 | } |
| 34 | |
| 35 | // MockProvider is a provider.Provider whose Stream returns scripted |
| 36 | // responses, one Turn per call. It records every request it receives so |
| 37 | // tests can inspect what was sent to the model (cache surface, tool |
| 38 | // schemas, message ordering). |
| 39 | // |
| 40 | // Usage: |
| 41 | // |
| 42 | // mp := NewMock("test-model", Turn{Text: "Hello"}).Record() |
| 43 | // agent := agent.New(mp, registry, session, opts, nil) |
| 44 | // agent.Run(ctx, "hi") |
| 45 | // |
| 46 | // for i, req := range mp.Requests() { |
| 47 | // fmt.Printf("turn %d: %d messages, %d tools\n", i+1, |
| 48 | // len(req.Messages), len(req.Tools)) |
| 49 | // } |
| 50 | type MockProvider struct { |
| 51 | mu sync.Mutex |
| 52 | name string |
| 53 | script []Turn |
| 54 | seen int |
| 55 | reqs []provider.Request |
| 56 | } |
| 57 | |
| 58 | // NewMock creates a MockProvider. The turns argument is the script; each |
| 59 | // Stream call consumes one Turn. Extra calls after the script are exhausted |
| 60 | // return an error. Use Append or SetScript to add more turns later. |
| 61 | func NewMock(name string, turns ...Turn) *MockProvider { |
| 62 | return &MockProvider{name: name, script: turns} |
| 63 | } |
| 64 | |
| 65 | // Name returns the provider instance name. |
| 66 | func (p *MockProvider) Name() string { return p.name } |
| 67 | |
| 68 | // Stream replays the next scripted turn. It records the request, then |
| 69 | // sends chunks in order (reasoning → text → tool calls → usage → done). |
| 70 | // If the Turn has StreamError set it is returned immediately. |
| 71 | func (p *MockProvider) Stream(ctx context.Context, req provider.Request) (<-chan provider.Chunk, error) { |
| 72 | if err := ctx.Err(); err != nil { |
| 73 | return nil, err |
| 74 | } |
| 75 | |
| 76 | p.mu.Lock() |
| 77 | p.reqs = append(p.reqs, req) |
| 78 | if p.seen >= len(p.script) { |
| 79 | p.mu.Unlock() |
| 80 | return nil, fmt.Errorf("MockProvider[%s]: no scripted turn %d (have %d turns)", p.name, p.seen, len(p.script)) |
| 81 | } |
| 82 | t := p.script[p.seen] |
| 83 | p.seen++ |
| 84 | p.mu.Unlock() |
| 85 | if t.StreamError != nil { |
| 86 | return nil, t.StreamError |
| 87 | } |
| 88 | |
| 89 | var chunks []provider.Chunk |
| 90 | if len(t.Chunks) > 0 { |
| 91 | chunks = append(chunks, t.Chunks...) |
| 92 | } else { |
| 93 | if t.Reasoning != "" { |
| 94 | chunks = append(chunks, provider.Chunk{Type: provider.ChunkReasoning, Text: t.Reasoning}) |
| 95 | } |
| 96 | if t.Text != "" { |
| 97 | chunks = append(chunks, provider.Chunk{Type: provider.ChunkText, Text: t.Text}) |
| 98 | } |
| 99 | for i := range t.ToolCalls { |
| 100 | tc := t.ToolCalls[i] |
| 101 | chunks = append(chunks, provider.Chunk{Type: provider.ChunkToolCall, ToolCall: &tc}) |
| 102 | } |
| 103 | if t.Usage != nil { |
| 104 | chunks = append(chunks, provider.Chunk{Type: provider.ChunkUsage, Usage: t.Usage}) |
| 105 | } |
| 106 | if t.ChunkError != nil { |
| 107 | chunks = append(chunks, provider.Chunk{Type: provider.ChunkError, Err: t.ChunkError}) |
| 108 | } else { |
| 109 | chunks = append(chunks, provider.Chunk{Type: provider.ChunkDone}) |
| 110 | } |
| 111 | } |
| 112 | |
| 113 | ch := make(chan provider.Chunk) |
| 114 | go func() { |
| 115 | defer close(ch) |
| 116 | for _, c := range chunks { |
| 117 | // Check before every send so an already-observed cancellation is |
| 118 | // never masked by a select that also has a ready receiver. |
| 119 | if err := ctx.Err(); err != nil { |
| 120 | ch <- provider.Chunk{Type: provider.ChunkError, Err: err} |
| 121 | return |
| 122 | } |
| 123 | select { |
| 124 | case <-ctx.Done(): |
| 125 | ch <- provider.Chunk{Type: provider.ChunkError, Err: ctx.Err()} |
| 126 | return |
| 127 | case ch <- c: |
| 128 | } |
| 129 | } |
| 130 | }() |
| 131 | return ch, nil |
| 132 | } |
| 133 | |
| 134 | // Requests returns all recorded requests in call order. Safe to call from |
| 135 | // any goroutine after the run loop finishes. |
| 136 | func (p *MockProvider) Requests() []provider.Request { |
| 137 | p.mu.Lock() |
| 138 | defer p.mu.Unlock() |
| 139 | out := make([]provider.Request, len(p.reqs)) |
| 140 | copy(out, p.reqs) |
| 141 | return out |
| 142 | } |
| 143 | |
| 144 | // LastRequest returns the most recent request, or nil if none. |
| 145 | func (p *MockProvider) LastRequest() *provider.Request { |
| 146 | p.mu.Lock() |
| 147 | defer p.mu.Unlock() |
| 148 | if len(p.reqs) == 0 { |
| 149 | return nil |
| 150 | } |
| 151 | r := p.reqs[len(p.reqs)-1] |
| 152 | return &r |
| 153 | } |
| 154 | |
| 155 | // MessageCount is a shortcut for len(Requests()). |
| 156 | func (p *MockProvider) CallCount() int { |
| 157 | p.mu.Lock() |
| 158 | defer p.mu.Unlock() |
| 159 | return p.seen |
| 160 | } |
| 161 | |
| 162 | // SetScript replaces the script and resets the call counter. |
| 163 | func (p *MockProvider) SetScript(turns ...Turn) { |
| 164 | p.mu.Lock() |
| 165 | defer p.mu.Unlock() |
| 166 | p.script = turns |
| 167 | p.seen = 0 |
| 168 | } |
| 169 | |
| 170 | // Append adds turns to the existing script. |
| 171 | func (p *MockProvider) Append(turns ...Turn) { |
| 172 | p.mu.Lock() |
| 173 | defer p.mu.Unlock() |
| 174 | p.script = append(p.script, turns...) |
| 175 | } |
| 176 | |
| 177 | // Reset clears recorded requests and the call counter without changing the script. |
| 178 | func (p *MockProvider) Reset() { |
| 179 | p.mu.Lock() |
| 180 | defer p.mu.Unlock() |
| 181 | p.reqs = nil |
| 182 | p.seen = 0 |
| 183 | } |
| 184 | |
| 185 | // UsageTurn is a convenience: a Turn whose text is empty but usage is set. |
| 186 | // Useful for simulating a tool-call round-trip where the final model response |
| 187 | // that round is tested later. |
| 188 | func UsageTurn(hit, miss, completion int) Turn { |
| 189 | return Turn{ |
| 190 | Usage: &provider.Usage{ |
| 191 | CacheHitTokens: hit, |
| 192 | CacheMissTokens: miss, |
| 193 | CompletionTokens: completion, |
| 194 | PromptTokens: hit + miss, |
| 195 | TotalTokens: hit + miss + completion, |
| 196 | }, |
| 197 | } |
| 198 | } |
| 199 | |
| 200 | // ErrorTurn is a convenience: a Turn that immediately returns the given error. |
| 201 | func ErrorTurn(err error) Turn { |
| 202 | return Turn{StreamError: err} |
| 203 | } |
| 204 | |
| 205 | var _ provider.Provider = (*MockProvider)(nil) |
| 206 |