print_test.go RawHistoryFolder

1package modes
2
3import (
4	"bytes"
5	"context"
6	"testing"
7
8	"github.com/patriceckhart/zot/packages/core"
9	"github.com/patriceckhart/zot/packages/provider"
10)
11
12type printTestClient struct{}
13
14func (printTestClient) Name() string { return "test" }
15
16func (printTestClient) Stream(context.Context, provider.Request) (<-chan provider.Event, error) {
17	events := make(chan provider.Event, 3)
18	events <- provider.EventStart{Model: "test-model", Provider: "test"}
19	events <- provider.EventUsage{Usage: provider.Usage{
20		InputTokens:          12,
21		OutputTokens:         9,
22		ReasoningTokens:      6,
23		ReasoningTokensKnown: true,
24	}}
25	events <- provider.EventDone{
26		Stop: provider.StopEnd,
27		Message: provider.Message{Role: provider.RoleAssistant, Content: []provider.Content{
28			provider.TextBlock{Text: "answer"},
29		}},
30	}
31	close(events)
32	return events, nil
33}
34
35func TestRunPrintReturnsInvocationUsage(t *testing.T) {
36	agent := core.NewAgent(printTestClient{}, "test-model", "", core.Registry{})
37	agent.SeedCost(provider.Usage{InputTokens: 100, OutputTokens: 50})
38	var output bytes.Buffer
39
40	usage, err := RunPrint(context.Background(), agent, "question", nil, &output)
41	if err != nil {
42		t.Fatal(err)
43	}
44	if output.String() != "answer\n" {
45		t.Fatalf("output = %q", output.String())
46	}
47	if usage.InputTokens != 12 || usage.OutputTokens != 9 || usage.ReasoningTokens != 6 || !usage.ReasoningTokensKnown {
48		t.Fatalf("usage = %+v", usage)
49	}
50}