go-deepseek/sse_test.go

package deepseek

import (
	"context"
	"errors"
	"io"
	"reflect"
	"strings"
	"testing"
)

type failingReader struct {
	err error
}

func (r failingReader) Read([]byte) (int, error) { return 0, r.err }

func TestSSEReaderNext(t *testing.T) {
	input := strings.Join([]string{
		`: keepalive`,
		`id: event-id`,
		`event: update`,
		`retry: 1000`,
		`data: first`,
		`data: second`,
		`unknown: ignored`,
		``,
	}, "\r\n")

	event, err := NewSSEReader(strings.NewReader(input)).Next(context.Background())
	if err != nil {
		t.Fatalf("Next() error = %v", err)
	}
	want := SSEEvent{ID: "event-id", Event: "update", Retry: "1000", Data: []byte("first\nsecond")}
	if !reflect.DeepEqual(event, want) {
		t.Fatalf("event = %#v, want %#v", event, want)
	}
}

func TestSSEReaderNextReturnsFinalEventWithoutBlankLine(t *testing.T) {
	event, err := NewSSEReader(strings.NewReader("event: close\ndata: {}")).Next(context.Background())
	if err != nil {
		t.Fatalf("Next() error = %v", err)
	}
	if event.Event != "close" || string(event.Data) != "{}" {
		t.Fatalf("event = %#v", event)
	}
}

func TestSSEReaderNextEOF(t *testing.T) {
	_, err := NewSSEReader(strings.NewReader(": comment without event")).Next(context.Background())
	if !errors.Is(err, io.EOF) {
		t.Fatalf("Next() error = %v, want io.EOF", err)
	}
}

func TestSSEReaderNextContextCanceled(t *testing.T) {
	ctx, cancel := context.WithCancel(context.Background())
	cancel()
	_, err := NewSSEReader(strings.NewReader("data: ignored\n\n")).Next(ctx)
	if !errors.Is(err, context.Canceled) {
		t.Fatalf("Next() error = %v, want context canceled", err)
	}
}

func TestSSEReaderNextReadError(t *testing.T) {
	wantErr := errors.New("read failed")
	_, err := NewSSEReader(failingReader{err: wantErr}).Next(context.Background())
	if !errors.Is(err, wantErr) {
		t.Fatalf("Next() error = %v, want %v", err, wantErr)
	}
}

func TestStreamStateApplyMetadataEvents(t *testing.T) {
	state := &StreamState{}
	events := []SSEEvent{
		{Event: "ready", Data: []byte(`{"request_message_id":11,"response_message_id":12,"model_type":"default"}`)},
		{Event: "update_session", Data: []byte(`{"updated_at":123.5}`)},
		{Event: "title", Data: []byte(`{"content":"Chat title"}`)},
		{Event: "finish"},
		{Event: "close"},
	}
	for _, event := range events {
		if err := state.Apply(event); err != nil {
			t.Fatalf("Apply(%q) error = %v", event.Event, err)
		}
	}
	if state.RequestMessageID != 11 || state.ResponseMessageID != 12 || state.SessionUpdated != 123.5 || state.Title != "Chat title" || !state.Finished || !state.Closed {
		t.Fatalf("state = %#v", state)
	}
}

func TestStreamStateApplySnapshotAndPatches(t *testing.T) {
	thinking := "initial thought"
	state := &StreamState{}
	snapshot := `{"v":{"response":{"message_id":7,"status":"PENDING","quasi_status":"STREAMING","accumulated_token_usage":3,"thinking_content":"` + thinking + `","fragments":[{"type":"RESPONSE","content":"initial"}]}}}`
	if err := state.Apply(SSEEvent{Data: []byte(snapshot)}); err != nil {
		t.Fatalf("apply snapshot: %v", err)
	}
	if state.ResponseMessageID != 7 || state.Status != "STREAMING" || state.TokenUsage != 3 || state.Content.String() != "initial" || state.ThinkingContent.String() != thinking {
		t.Fatalf("snapshot state = %#v", state)
	}

	patches := []string{
		`{"p":"response/content","o":"APPEND","v":" one"}`,
		`{"v":" two"}`,
		`{"p":"response/thinking_content","o":"SET","v":"replaced"}`,
		`{"p":"response","o":"BATCH","v":[{"p":"accumulated_token_usage","o":"SET","v":9},{"p":"quasi_status","o":"SET","v":"FINISHED"}]}`,
	}
	for _, patch := range patches {
		if err := state.Apply(SSEEvent{Data: []byte(patch)}); err != nil {
			t.Fatalf("Apply(%s) error = %v", patch, err)
		}
	}
	if state.Content.String() != "initial one two" {
		t.Fatalf("content = %q", state.Content.String())
	}
	if state.ThinkingContent.String() != "replaced" || state.TokenUsage != 9 || state.Status != "FINISHED" || !state.Finished {
		t.Fatalf("patched state = %#v", state)
	}
}

func TestStreamStateApplyErrors(t *testing.T) {
	tests := []struct {
		name  string
		event SSEEvent
		want  string
	}{
		{name: "invalid patch JSON", event: SSEEvent{Data: []byte(`{`)}, want: "decode stream patch"},
		{name: "invalid ready", event: SSEEvent{Event: "ready", Data: []byte(`{`)}, want: "decode ready"},
		{name: "invalid string", event: SSEEvent{Data: []byte(`{"p":"response/content","o":"SET","v":1}`)}, want: "decode string patch"},
		{name: "unsupported operation", event: SSEEvent{Data: []byte(`{"p":"response/content","o":"DELETE","v":"x"}`)}, want: "unsupported string patch operation"},
		{name: "invalid batch", event: SSEEvent{Data: []byte(`{"p":"response","o":"BATCH","v":{}}`)}, want: "decode batch"},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			err := (&StreamState{}).Apply(tt.event)
			if err == nil || !strings.Contains(err.Error(), tt.want) {
				t.Fatalf("Apply() error = %v, want containing %q", err, tt.want)
			}
		})
	}
}