homelab-frontend/internal/proxy/streaming_test.go

package proxy

import (
	"bufio"
	"fmt"
	"io"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
	"time"

	"forgejo.riotpiao.com/rock/homelab-frontend/internal/config"
)

// TestSSEUnbuffered verifies that SSE events stream to the client without buffering.
func TestSSEUnbuffered(t *testing.T) {
	// Upstream that emits SSE events with gaps
	sseEvents := []string{"event1", "event2", "event3", "event4", "event5"}
	eventGap := 20 * time.Millisecond

	upstreamServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "text/event-stream")
		w.Header().Set("Cache-Control", "no-cache")
		w.Header().Set("Connection", "keep-alive")
		w.WriteHeader(http.StatusOK)

		rc := http.NewResponseController(w)
		for i, event := range sseEvents {
			if i > 0 {
				time.Sleep(eventGap)
			}
			fmt.Fprintf(w, "data: %s\n\n", event)
			if err := rc.Flush(); err != nil {
				return
			}
		}
		fmt.Fprintf(w, "data: [DONE]\n\n")
		_ = rc.Flush()
	}))
	defer upstreamServer.Close()

	upstreamAddr := strings.TrimPrefix(upstreamServer.URL, "http://")

	cfg := &config.Config{
		Routes: map[string]*config.Route{
			"sse-route": {
				Name: "sse-route",
				Upstream: config.Upstream{
					Address:        upstreamAddr,
					ConnectTimeout: 5 * time.Second,
					ReadTimeout:    10 * time.Second,
					WriteTimeout:   5 * time.Second,
					MaxBodySize:    1024 * 1024,
					AuthRequired:   false,
				},
			},
		},
	}

	handler := New(cfg)
	defer handler.Close()

	server := httptest.NewServer(handler)
	defer server.Close()

	// Connect to the proxy
	resp, err := http.Get(server.URL + "/sse")
	if err != nil {
		t.Fatalf("request failed: %v", err)
	}
	defer resp.Body.Close()

	// Verify headers
	if resp.Header.Get("Content-Type") != "text/event-stream" {
		t.Errorf("expected Content-Type: text/event-stream, got %s", resp.Header.Get("Content-Type"))
	}
	if resp.Header.Get("Cache-Control") != "no-cache" {
		t.Errorf("expected Cache-Control: no-cache, got %s", resp.Header.Get("Cache-Control"))
	}

	// Read events and measure timing
	reader := bufio.NewReader(resp.Body)
	eventTimes := make([]time.Time, 0, len(sseEvents))
	observedEvents := make([]string, 0, len(sseEvents))

	for {
		line, err := reader.ReadString('\n')
		if err != nil {
			if err == io.EOF {
				break
			}
			t.Fatalf("read failed: %v", err)
		}

		line = strings.TrimSpace(line)
		if strings.HasPrefix(line, "data: ") {
			event := strings.TrimPrefix(line, "data: ")
			if event == "[DONE]" {
				break
			}
			eventTimes = append(eventTimes, time.Now())
			observedEvents = append(observedEvents, event)
		}
	}

	// Verify we got all events
	if len(observedEvents) != len(sseEvents) {
		t.Errorf("expected %d events, got %d", len(sseEvents), len(observedEvents))
	}

	// Verify events match
	for i, expected := range sseEvents {
		if i < len(observedEvents) && observedEvents[i] != expected {
			t.Errorf("event %d: expected %s, got %s", i, expected, observedEvents[i])
		}
	}

	// Verify timing gaps between events are reasonable
	// The gaps should be approximately eventGap (allowing for some overhead)
	for i := 1; i < len(eventTimes); i++ {
		gap := eventTimes[i].Sub(eventTimes[i-1])
		minGap := eventGap * 80 / 100  // Allow 20% tolerance
		maxGap := eventGap * 300 / 100 // Allow up to 3x the expected gap
		if gap < minGap || gap > maxGap {
			t.Logf("event gap %d: %.1fms (expected ~%.1fms)", i, gap.Seconds()*1000, eventGap.Seconds()*1000)
		}
	}
}

// TestChunkedUnbuffered verifies that chunked responses stream without buffering.
func TestChunkedUnbuffered(t *testing.T) {
	chunks := []string{"chunk1\n", "chunk2\n", "chunk3\n"}
	chunkGap := 20 * time.Millisecond

	upstreamServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "text/plain")
		w.WriteHeader(http.StatusOK)

		rc := http.NewResponseController(w)
		for i, chunk := range chunks {
			if i > 0 {
				time.Sleep(chunkGap)
			}
			fmt.Fprint(w, chunk)
			if err := rc.Flush(); err != nil {
				return
			}
		}
	}))
	defer upstreamServer.Close()

	upstreamAddr := strings.TrimPrefix(upstreamServer.URL, "http://")

	cfg := &config.Config{
		Routes: map[string]*config.Route{
			"chunked-route": {
				Name: "chunked-route",
				Upstream: config.Upstream{
					Address:        upstreamAddr,
					ConnectTimeout: 5 * time.Second,
					ReadTimeout:    10 * time.Second,
					WriteTimeout:   5 * time.Second,
					MaxBodySize:    1024 * 1024,
					AuthRequired:   false,
				},
			},
		},
	}

	handler := New(cfg)
	defer handler.Close()

	server := httptest.NewServer(handler)
	defer server.Close()

	resp, err := http.Get(server.URL + "/chunked")
	if err != nil {
		t.Fatalf("request failed: %v", err)
	}
	defer resp.Body.Close()

	// Read chunks and verify they arrive before all are sent
	reader := bufio.NewReader(resp.Body)
	receivedChunks := make([]string, 0, len(chunks))

	for {
		chunk := make([]byte, 0, 1024)
		for {
			b, err := reader.ReadByte()
			if err != nil {
				if err == io.EOF {
					break
				}
				t.Fatalf("read failed: %v", err)
			}
			chunk = append(chunk, b)
			if b == '\n' {
				break
			}
		}

		if len(chunk) > 0 {
			receivedChunks = append(receivedChunks, string(chunk))
		}

		if len(receivedChunks) >= len(chunks) {
			break
		}
	}

	// Verify chunks match
	if len(receivedChunks) != len(chunks) {
		t.Errorf("expected %d chunks, got %d", len(chunks), len(receivedChunks))
	}

	for i, expected := range chunks {
		if i < len(receivedChunks) && strings.TrimSpace(receivedChunks[i]) != strings.TrimSpace(expected) {
			t.Errorf("chunk %d: expected %s, got %s", i, strings.TrimSpace(expected), strings.TrimSpace(receivedChunks[i]))
		}
	}
}

// TestHeadersBeforeBody verifies that response headers reach the client before the body.
func TestHeadersBeforeBody(t *testing.T) {
	headersSent := make(chan bool, 1)

	upstreamServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("X-Custom-Header", "test-value")
		w.Header().Set("Content-Type", "text/plain")
		w.WriteHeader(http.StatusOK)
		headersSent <- true

		// Simulate slow body send
		time.Sleep(100 * time.Millisecond)
		fmt.Fprint(w, "body content")
	}))
	defer upstreamServer.Close()

	upstreamAddr := strings.TrimPrefix(upstreamServer.URL, "http://")

	cfg := &config.Config{
		Routes: map[string]*config.Route{
			"headers-route": {
				Name: "headers-route",
				Upstream: config.Upstream{
					Address:        upstreamAddr,
					ConnectTimeout: 5 * time.Second,
					ReadTimeout:    5 * time.Second,
					WriteTimeout:   5 * time.Second,
					MaxBodySize:    1024 * 1024,
					AuthRequired:   false,
				},
			},
		},
	}

	handler := New(cfg)
	defer handler.Close()

	server := httptest.NewServer(handler)
	defer server.Close()

	resp, err := http.Get(server.URL + "/headers")
	if err != nil {
		t.Fatalf("request failed: %v", err)
	}

	// Headers should be immediately available
	if resp.Header.Get("X-Custom-Header") != "test-value" {
		t.Errorf("custom header not received before body")
	}

	resp.Body.Close()
}

// TestResponseHeadersPassThrough verifies that various response headers survive the proxy.
func TestResponseHeadersPassThrough(t *testing.T) {
	testHeaders := map[string]string{
		"Content-Type":  "application/json",
		"Cache-Control": "no-cache, no-store",
		"X-Custom":      "custom-value",
	}

	upstreamServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		for k, v := range testHeaders {
			w.Header().Set(k, v)
		}
		w.WriteHeader(http.StatusOK)
		fmt.Fprint(w, "test")
	}))
	defer upstreamServer.Close()

	upstreamAddr := strings.TrimPrefix(upstreamServer.URL, "http://")

	cfg := &config.Config{
		Routes: map[string]*config.Route{
			"headers-route": {
				Name: "headers-route",
				Upstream: config.Upstream{
					Address:        upstreamAddr,
					ConnectTimeout: 5 * time.Second,
					ReadTimeout:    5 * time.Second,
					WriteTimeout:   5 * time.Second,
					MaxBodySize:    1024 * 1024,
					AuthRequired:   false,
				},
			},
		},
	}

	handler := New(cfg)
	defer handler.Close()

	server := httptest.NewServer(handler)
	defer server.Close()

	resp, err := http.Get(server.URL + "/test")
	if err != nil {
		t.Fatalf("request failed: %v", err)
	}
	defer resp.Body.Close()

	for k, v := range testHeaders {
		if resp.Header.Get(k) != v {
			t.Errorf("header %s: expected %s, got %s", k, v, resp.Header.Get(k))
		}
	}
}

// TestDONESentinel verifies that the [DONE] sentinel reaches the client.
func TestDONESentinel(t *testing.T) {
	upstreamServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "text/event-stream")
		w.WriteHeader(http.StatusOK)

		rc := http.NewResponseController(w)
		fmt.Fprint(w, "data: token1\n\n")
		_ = rc.Flush()
		fmt.Fprint(w, "data: [DONE]\n\n")
		_ = rc.Flush()
	}))
	defer upstreamServer.Close()

	upstreamAddr := strings.TrimPrefix(upstreamServer.URL, "http://")

	cfg := &config.Config{
		Routes: map[string]*config.Route{
			"sse-route": {
				Name: "sse-route",
				Upstream: config.Upstream{
					Address:        upstreamAddr,
					ConnectTimeout: 5 * time.Second,
					ReadTimeout:    5 * time.Second,
					WriteTimeout:   5 * time.Second,
					MaxBodySize:    1024 * 1024,
					AuthRequired:   false,
				},
			},
		},
	}

	handler := New(cfg)
	defer handler.Close()

	server := httptest.NewServer(handler)
	defer server.Close()

	resp, err := http.Get(server.URL + "/sse")
	if err != nil {
		t.Fatalf("request failed: %v", err)
	}
	defer resp.Body.Close()

	reader := bufio.NewReader(resp.Body)
	foundDONE := false

	for {
		line, err := reader.ReadString('\n')
		if err != nil {
			if err == io.EOF {
				break
			}
			t.Fatalf("read failed: %v", err)
		}

		line = strings.TrimSpace(line)
		if strings.Contains(line, "[DONE]") {
			foundDONE = true
			break
		}
	}

	if !foundDONE {
		t.Errorf("expected [DONE] sentinel, not found")
	}
}

// TestNoFullBuffering verifies that the response is not fully buffered in memory.
func TestNoFullBuffering(t *testing.T) {
	// Create a large response that would be problematic if fully buffered
	chunkCount := 10
	chunkSize := 100000
	largeData := strings.Repeat("x", chunkCount*chunkSize)

	upstreamServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "text/plain")
		w.WriteHeader(http.StatusOK)

		rc := http.NewResponseController(w)
		// Send data in chunks with gaps to ensure streaming
		for i := 0; i < chunkCount; i++ {
			chunk := largeData[i*chunkSize : (i+1)*chunkSize]
			fmt.Fprint(w, chunk)
			if err := rc.Flush(); err != nil {
				return
			}
			if i < chunkCount-1 {
				time.Sleep(10 * time.Millisecond)
			}
		}
	}))
	defer upstreamServer.Close()

	upstreamAddr := strings.TrimPrefix(upstreamServer.URL, "http://")

	cfg := &config.Config{
		Routes: map[string]*config.Route{
			"large-route": {
				Name: "large-route",
				Upstream: config.Upstream{
					Address:        upstreamAddr,
					ConnectTimeout: 5 * time.Second,
					ReadTimeout:    10 * time.Second,
					WriteTimeout:   5 * time.Second,
					MaxBodySize:    100 * 1024 * 1024,
					AuthRequired:   false,
				},
			},
		},
	}

	handler := New(cfg)
	defer handler.Close()

	server := httptest.NewServer(handler)
	defer server.Close()

	resp, err := http.Get(server.URL + "/large")
	if err != nil {
		t.Fatalf("request failed: %v", err)
	}
	defer resp.Body.Close()

	// Read the response in chunks to verify streaming
	totalRead := 0
	readChunkSize := 8192
	for {
		buf := make([]byte, readChunkSize)
		n, err := resp.Body.Read(buf)
		if n > 0 {
			totalRead += n
		}
		if err != nil {
			if err == io.EOF {
				break
			}
			t.Fatalf("read failed: %v", err)
		}
	}

	if totalRead != len(largeData) {
		t.Errorf("expected to read %d bytes, got %d", len(largeData), totalRead)
	}
}