1package modes
2
3import (
4 "bytes"
5 "context"
6 "testing"
7
8 "github.com/bnema/zut/packages/core"
9 "github.com/bnema/zut/packages/provider"
10)
11
12type streamTestClient struct {
13 events []provider.Event
14}
15
16func (c *streamTestClient) Name() string { return "stream-test" }
17
18func (c *streamTestClient) Stream(ctx context.Context, req provider.Request) (<-chan provider.Event, error) {
19 out := make(chan provider.Event, len(c.events))
20 go func() {
21 defer close(out)
22 for _, ev := range c.events {
23 select {
24 case <-ctx.Done():
25 return
26 case out <- ev:
27 }
28 }
29 }()
30 return out, nil
31}
32
33func TestRunStreamWritesDeltasLive(t *testing.T) {
34 client := &streamTestClient{events: []provider.Event{
35 provider.EventTextDelta{Delta: "Hel"},
36 provider.EventTextDelta{Delta: "lo"},
37 provider.EventDone{
38 Stop: provider.StopEnd,
39 Message: provider.Message{
40 Role: provider.RoleAssistant,
41 Content: []provider.Content{provider.TextBlock{Text: "Hello"}},
42 },
43 },
44 }}
45 ag := core.NewAgent(client, "test-model", "", nil)
46 var out, diag bytes.Buffer
47 if err := RunStreamWithDiag(context.Background(), ag, "hi", nil, &out, &diag); err != nil {
48 t.Fatal(err)
49 }
50 if got := out.String(); got != "Hello\n" {
51 t.Fatalf("stdout = %q, want Hello\\n", got)
52 }
53}
54
55func TestRunStreamFallbackWithoutDeltas(t *testing.T) {
56 client := &streamTestClient{events: []provider.Event{
57 provider.EventDone{
58 Stop: provider.StopEnd,
59 Message: provider.Message{
60 Role: provider.RoleAssistant,
61 Content: []provider.Content{provider.TextBlock{Text: "done"}},
62 },
63 },
64 }}
65 ag := core.NewAgent(client, "test-model", "", nil)
66 var out bytes.Buffer
67 if err := RunStreamWithDiag(context.Background(), ag, "hi", nil, &out, ioDiscard{}); err != nil {
68 t.Fatal(err)
69 }
70 if got := out.String(); got != "done\n" {
71 t.Fatalf("stdout = %q, want done\\n", got)
72 }
73}
74
75type ioDiscard struct{}
76
77func (ioDiscard) Write(p []byte) (int, error) { return len(p), nil }