28 lines
976 B
Go
28 lines
976 B
Go
|
package main
|
||
|
|
|
||
|
|
import (
|
||
|
|
"bytes"
|
||
|
|
"testing"
|
||
|
|
)
|
||
|
|
|
||
|
|
func TestEvents(t *testing.T) {
|
||
|
|
finish := []byte("data: {\"choices\":[{\"delta\":{\"content\":\"answer\"},\"finish_reason\":\"stop\"}]}\n\n")
|
||
|
|
usage := []byte("data: {\"choices\":[],\"usage\":{\"total_tokens\":16}}\n\n")
|
||
|
|
done := []byte("data: [DONE]\n\n")
|
||
|
|
input := append(append(append([]byte(nil), finish...), usage...), done...)
|
||
|
|
for split := 0; split <= len(input); split++ {
|
||
|
|
var parser events
|
||
|
|
output := parser.feed(input[:split], false)
|
||
|
|
output = append(output, parser.feed(input[split:], true)...)
|
||
|
|
if !bytes.Equal(output, input) || !parser.done {
|
||
|
|
t.Fatalf("split %d: %q", split, output)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
for _, input := range [][]byte{done, []byte("data: {\"choices\":[{\"delta\":{\"role\":\"assistant\"}}]}\n\ndata: [DONE]\n\n"), nil} {
|
||
|
|
var parser events
|
||
|
|
output := parser.feed(input, true)
|
||
|
|
if !bytes.Contains(output, streamError) || bytes.Contains(output, done) {
|
||
|
|
t.Fatalf("failed stream: %q", output)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|