agentic-coding-workflow/internal/chat/sse_test.go

package chat

import (
	"context"
	"encoding/json"
	"fmt"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
	"time"

	"github.com/go-chi/chi/v5"
	"github.com/google/uuid"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"

	"github.com/yan1h/agent-coding-workflow/internal/infra/errs"
	"github.com/yan1h/agent-coding-workflow/internal/transport/http/middleware"
)

// ===== SSE-capable response writer =====

// flushRecorder is an httptest.ResponseRecorder that also implements http.Flusher.
type flushRecorder struct {
	*httptest.ResponseRecorder
	flushed int
}

func newFlushRecorder() *flushRecorder {
	return &flushRecorder{ResponseRecorder: httptest.NewRecorder()}
}

func (f *flushRecorder) Flush() {
	f.flushed++
}

// ===== Fake MessageService for SSE tests =====

// sseMessageService allows controlling the Stream channel.
type sseMessageService struct {
	streamCh  chan SSEEvent
	streamErr error
}

func (s *sseMessageService) Send(_ context.Context, _, _ uuid.UUID, _ string, _ []uuid.UUID) (int64, int64, error) {
	return 0, 0, nil
}
func (s *sseMessageService) Stream(_ context.Context, _ uuid.UUID, _ uuid.UUID, _ int64) (<-chan SSEEvent, error) {
	if s.streamErr != nil {
		return nil, s.streamErr
	}
	return s.streamCh, nil
}
func (s *sseMessageService) Cancel(_ context.Context, _ uuid.UUID, _ int64) error { return nil }
func (s *sseMessageService) Retry(_ context.Context, _ uuid.UUID, _ int64) (int64, uuid.UUID, error) {
	return 0, uuid.Nil, nil
}
func (s *sseMessageService) ListMessages(_ context.Context, _, _ uuid.UUID, _ *int64, _ int) ([]Message, error) {
	return nil, nil
}

// ===== SSE helper unit tests =====

// TestSSE_WritesEventInProperFormat verifies that writeSSE produces the
// canonical id/event/data/\n\n format.
func TestSSE_WritesEventInProperFormat(t *testing.T) {
	w := newFlushRecorder()
	ev := SSEEvent{
		ID:    42,
		Event: "text_delta",
		Data:  map[string]string{"delta": "hello"},
	}
	writeSSE(w, w, ev)

	body := w.Body.String()
	assert.True(t, strings.HasPrefix(body, "id: 42\n"), "must start with id line")
	assert.Contains(t, body, "event: text_delta\n")
	assert.Contains(t, body, "data: ")
	assert.True(t, strings.HasSuffix(body, "\n\n"), "must end with double newline")
	assert.Equal(t, 1, w.flushed)

	// Verify the data line is valid JSON.
	lines := strings.Split(strings.TrimRight(body, "\n"), "\n")
	var dataLine string
	for _, l := range lines {
		if strings.HasPrefix(l, "data: ") {
			dataLine = strings.TrimPrefix(l, "data: ")
		}
	}
	require.NotEmpty(t, dataLine)
	var payload map[string]string
	require.NoError(t, json.Unmarshal([]byte(dataLine), &payload))
	assert.Equal(t, "hello", payload["delta"])
}

// TestSSE_WriteSSEError verifies that writeSSEError produces an error event
// with the expected code/message fields from an AppError.
func TestSSE_WriteSSEError(t *testing.T) {
	w := newFlushRecorder()
	appErr := errs.New(errs.CodeChatConversationNotFound, "conversation gone")
	writeSSEError(w, w, appErr)

	body := w.Body.String()
	assert.Contains(t, body, "event: error\n")
	assert.True(t, strings.HasSuffix(body, "\n\n"))
	assert.Equal(t, 1, w.flushed)

	lines := strings.Split(body, "\n")
	var dataLine string
	for _, l := range lines {
		if strings.HasPrefix(l, "data: ") {
			dataLine = strings.TrimPrefix(l, "data: ")
		}
	}
	require.NotEmpty(t, dataLine)
	var payload map[string]string
	require.NoError(t, json.Unmarshal([]byte(dataLine), &payload))
	assert.Equal(t, string(errs.CodeChatConversationNotFound), payload["code"])
	assert.Equal(t, "conversation gone", payload["message"])
}

// TestStreamMessages_ContentTypeHeader verifies that the SSE handler sets
// the required Content-Type and X-Accel-Buffering headers.
func TestStreamMessages_ContentTypeHeader(t *testing.T) {
	uid := uuid.New()
	convID := uuid.New()

	ch := make(chan SSEEvent)
	close(ch) // immediately closed — no events, handler returns quickly.

	msgSvc := &sseMessageService{streamCh: ch}
	resolver := &fakeSessionResolver{uid: uid}
	upload := NewUploader(newFakeRepo(), nil, nil, 1<<20)
	h := NewHandler(&fakeConvSvc{}, msgSvc, &fakeTplSvc{}, &fakeEpSvc{}, &fakeUsageSvc{}, upload, resolver, newFakeAdminLookup())

	r := chi.NewRouter()
	r.Use(middleware.RequestID)
	h.Mount(r)

	req := httptest.NewRequest(http.MethodGet, fmt.Sprintf("/api/v1/conversations/%s/stream", convID), nil)
	req.Header.Set("Authorization", "Bearer test-token")
	w := newFlushRecorder()
	r.ServeHTTP(w, req)

	assert.Equal(t, "text/event-stream", w.Header().Get("Content-Type"))
	assert.Equal(t, "no-cache", w.Header().Get("Cache-Control"))
	assert.Equal(t, "no", w.Header().Get("X-Accel-Buffering"))
	assert.Equal(t, http.StatusOK, w.Code)
}

// TestStreamMessages_StopsWhenClientClosesContext verifies that when the
// request context is cancelled mid-stream, the handler stops consuming events
// and the channel is not blocked.
func TestStreamMessages_StopsWhenClientClosesContext(t *testing.T) {
	uid := uuid.New()
	convID := uuid.New()

	// Channel that never closes on its own — cancelling the ctx should stop the handler.
	ch := make(chan SSEEvent, 1)

	msgSvc := &sseMessageService{streamCh: ch}
	resolver := &fakeSessionResolver{uid: uid}
	upload := NewUploader(newFakeRepo(), nil, nil, 1<<20)
	h := NewHandler(&fakeConvSvc{}, msgSvc, &fakeTplSvc{}, &fakeEpSvc{}, &fakeUsageSvc{}, upload, resolver, newFakeAdminLookup())

	// Build a router identical to newTestHandler but without chi's context copying issues.
	r := chi.NewRouter()
	r.Use(middleware.RequestID)
	h.Mount(r)

	// Create a cancellable request.
	ctx, cancel := context.WithCancel(context.Background())
	req := httptest.NewRequest(http.MethodGet, fmt.Sprintf("/api/v1/conversations/%s/stream", convID), nil)
	req.Header.Set("Authorization", "Bearer test-token")
	req = req.WithContext(ctx)

	w := newFlushRecorder()

	done := make(chan struct{})
	go func() {
		defer close(done)
		r.ServeHTTP(w, req)
	}()

	// Send one event and give the handler a moment to drain it before cancelling,
	// so the body assertion is deterministic. The cancel itself is what we're
	// verifying — that the handler exits even though the channel never closes.
	ch <- SSEEvent{ID: 1, Event: "text_delta", Data: "hi"}
	require.Eventually(t, func() bool {
		return strings.Contains(w.Body.String(), "text_delta")
	}, time.Second, 10*time.Millisecond, "event was not flushed before cancel")
	cancel()

	// Handler must return promptly after context cancellation.
	select {
	case <-done:
		// success
	case <-waitTimeout(t, 2*time.Second):
		t.Fatal("handler did not stop after context cancellation")
	}

	assert.Contains(t, w.Body.String(), "text_delta")
}

// waitTimeout returns a channel that closes after the given duration.
func waitTimeout(t *testing.T, d time.Duration) <-chan time.Time {
	t.Helper()
	return time.After(d)
}