add a test for ErrTooManyEmptyStreamMessages in stream_reader (sashab…

…aranov#304)
vvatanabe · Jun 1, 2023 · 444c808 · 444c808
1 parent 15653fd
commit 444c808
Showing 1 changed file with 15 additions and 0 deletions.
diff --git a/stream_reader_test.go b/stream_reader_test.go
@@ -1,6 +1,8 @@
 package openai //nolint:testpackage // testing private field
 
 import (
+	"bufio"
+	"bytes"
 	"errors"
 	"testing"
 
@@ -36,3 +38,16 @@ func TestStreamReaderReturnsUnmarshalerErrors(t *testing.T) {
 		t.Fatalf("Did not return nil when unmarshaler failed: %v", respErr)
 	}
 }
+
+func TestStreamReaderReturnsErrTooManyEmptyStreamMessages(t *testing.T) {
+	stream := &streamReader[ChatCompletionStreamResponse]{
+		emptyMessagesLimit: 3,
+		reader:             bufio.NewReader(bytes.NewReader([]byte("\n\n\n\n"))),
+		errAccumulator:     utils.NewErrorAccumulator(),
+		unmarshaler:        &utils.JSONUnmarshaler{},
+	}
+	_, err := stream.Recv()
+	if !errors.Is(err, ErrTooManyEmptyStreamMessages) {
+		t.Fatalf("Did not return error when recv failed: %v", err)
+	}
+}