ds2api/internal/stream/engine_test.go

package stream

import (
	"context"
	"strings"
	"testing"

	"ds2api/internal/sse"
)

func TestConsumeSSEPrefersContextCancellationOverReadyParsedLines(t *testing.T) {
	ctx, cancel := context.WithCancel(context.Background())
	cancel()

	var finalized bool
	var contextDone bool
	var parsedCalled bool

	ConsumeSSE(ConsumeConfig{
		Context:           ctx,
		Body:              strings.NewReader("data: {\"p\":\"response/content\",\"v\":\"hello\"}\n\ndata: [DONE]\n"),
		ThinkingEnabled:   false,
		InitialType:       "text",
		KeepAliveInterval: 0,
	}, ConsumeHooks{
		OnParsed: func(_ sse.LineResult) ParsedDecision {
			parsedCalled = true
			return ParsedDecision{}
		},
		OnFinalize: func(_ StopReason, _ error) {
			finalized = true
		},
		OnContextDone: func() {
			contextDone = true
		},
	})

	if !contextDone {
		t.Fatal("expected OnContextDone to run for an already-cancelled context")
	}
	if finalized {
		t.Fatal("expected OnFinalize not to run after context cancellation wins")
	}
	if parsedCalled {
		t.Fatal("expected parsed lines not to be processed after context cancellation wins")
	}
}