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