ce7cb702218693d08ea6d83464d2457c9250027c

services/vertex-batch/stream_test.go

@@ -0,0 +1,28 @@
+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)
+		}
+	}
+}