| 1 | package provider |
| 2 | |
| 3 | import ( |
| 4 | "context" |
| 5 | "strings" |
| 6 | "testing" |
| 7 | ) |
| 8 | |
| 9 | type auxiliaryScript struct { |
| 10 | calls int |
| 11 | partial bool |
| 12 | failed bool |
| 13 | } |
| 14 | |
| 15 | func (*auxiliaryScript) Name() string { return "auxiliary-test" } |
| 16 | func (p *auxiliaryScript) Stream(context.Context, Request) (<-chan Chunk, error) { |
| 17 | p.calls++ |
| 18 | if p.failed { |
| 19 | return nil, &APIError{Status: 503} |
| 20 | } |
| 21 | ch := make(chan Chunk, 3) |
| 22 | ch <- Chunk{Type: ChunkText, Text: "candidate"} |
| 23 | ch <- Chunk{Type: ChunkUsage, Usage: &Usage{PromptTokens: 10, CompletionTokens: 1, TotalTokens: 11}} |
| 24 | if !p.partial { |
| 25 | ch <- Chunk{Type: ChunkDone} |
| 26 | } |
| 27 | close(ch) |
| 28 | return ch, nil |
| 29 | } |
| 30 | func TestAuxiliaryFailsWithoutRetryOrLeakingPartialText(t *testing.T) { |
| 31 | for _, mode := range []string{"complete", "partial", "upstream-error"} { |
| 32 | t.Run(mode, func(t *testing.T) { |
| 33 | p := &auxiliaryScript{partial: mode == "partial", failed: mode == "upstream-error"} |
| 34 | ch, err := StreamAuxiliary(t.Context(), p, Request{}) |
| 35 | if err != nil { |
| 36 | t.Fatal(err) |
| 37 | } |
| 38 | var text strings.Builder |
| 39 | var usage *Usage |
| 40 | var streamErr error |
| 41 | for c := range ch { |
| 42 | switch c.Type { |
| 43 | case ChunkText: |
| 44 | text.WriteString(c.Text) |
| 45 | case ChunkUsage: |
| 46 | usage = c.Usage |
| 47 | case ChunkError: |
| 48 | streamErr = c.Err |
| 49 | } |
| 50 | } |
| 51 | if p.calls != 1 || usage == nil || usage.RequestCount != 1 { |
| 52 | t.Fatalf("calls=%d usage=%+v", p.calls, usage) |
| 53 | } |
| 54 | if mode == "complete" { |
| 55 | if text.String() != "candidate" || streamErr != nil || usage.PromptTokens != 10 { |
| 56 | t.Fatalf("text=%q err=%v usage=%+v", text.String(), streamErr, usage) |
| 57 | } |
| 58 | } else if text.Len() != 0 || streamErr == nil { |
| 59 | t.Fatalf("partial text=%q err=%v", text.String(), streamErr) |
| 60 | } |
| 61 | }) |
| 62 | } |
| 63 | } |
| 64 |