homelab-frontend/internal/proxy/proxy_test.go

package proxy

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

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

func TestProxyBasic(t *testing.T) {
	// Start a stub upstream
	upstreamServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("X-Upstream-Header", "test-value")
		w.WriteHeader(http.StatusOK)
		fmt.Fprintf(w, "upstream response")
	}))
	defer upstreamServer.Close()

	// Extract host:port from upstream URL
	upstreamAddr := strings.TrimPrefix(upstreamServer.URL, "http://")

	// Create config with route to the stub
	cfg := &config.Config{
		Routes: map[string]*config.Route{
			"test-route": {
				Name: "test-route",
				Upstream: config.Upstream{
					Address:        upstreamAddr,
					PathRewrite:    "",
					ConnectTimeout: 5 * time.Second,
					ReadTimeout:    5 * time.Second,
					WriteTimeout:   5 * time.Second,
					MaxBodySize:    1024 * 1024,
					AuthRequired:   false,
				},
			},
		},
	}

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

	// Make a request through the proxy
	server := httptest.NewServer(handler)
	defer server.Close()

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

	// Verify status code
	if resp.StatusCode != http.StatusOK {
		t.Errorf("expected status 200, got %d", resp.StatusCode)
	}

	// Verify response header passed through
	if resp.Header.Get("X-Upstream-Header") != "test-value" {
		t.Errorf("upstream header not passed through")
	}

	// Verify response body
	body, _ := io.ReadAll(resp.Body)
	if string(body) != "upstream response" {
		t.Errorf("expected body 'upstream response', got %s", string(body))
	}
}

func TestProxyPathRewrite(t *testing.T) {
	requestedPath := ""
	upstreamServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		requestedPath = r.URL.Path
		w.WriteHeader(http.StatusOK)
	}))
	defer upstreamServer.Close()

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

	cfg := &config.Config{
		Routes: map[string]*config.Route{
			"rewrite-route": {
				Name: "rewrite-route",
				Upstream: config.Upstream{
					Address:        upstreamAddr,
					PathRewrite:    "/api/v2",
					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()

	// Not /v1/models: ServeHTTP serves that endpoint from config (task 2.5)
	// and returns before routing, so it never exercises PathRewrite.
	resp, err := http.Get(server.URL + "/some/path")
	if err != nil {
		t.Fatalf("request failed: %v", err)
	}
	resp.Body.Close()

	if requestedPath != "/api/v2" {
		t.Errorf("expected rewritten path /api/v2, got %s", requestedPath)
	}
}

func TestProxyConnectionReuse(t *testing.T) {
	upstreamServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusOK)
		fmt.Fprintf(w, "response")
	}))
	defer upstreamServer.Close()

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

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

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

	// Verify connection reuse by checking that the same transport is used

	// We can't easily count raw TCP connections in this test setup,
	// but we can verify that the transport is being reused by checking
	// that the same transport handles both requests
	route := handler.routes["test-route"]
	firstTransport := route.Transport

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

	// Make two sequential requests
	http.Get(server.URL + "/path1")
	http.Get(server.URL + "/path2")

	// Verify the same transport is used (connection reuse)
	if handler.transports[upstreamAddr] != firstTransport {
		t.Errorf("transport changed between requests")
	}

	// The transport should have been created once
	if len(handler.transports) != 1 {
		t.Errorf("expected 1 transport, got %d", len(handler.transports))
	}
}

func TestProxyNotFound(t *testing.T) {
	cfg := &config.Config{
		Routes: make(map[string]*config.Route),
	}

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

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

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

	if resp.StatusCode != http.StatusNotFound {
		t.Errorf("expected status 404, got %d", resp.StatusCode)
	}
}

func TestProxyUpstreamError(t *testing.T) {
	cfg := &config.Config{
		Routes: map[string]*config.Route{
			"bad-route": {
				Name: "bad-route",
				Upstream: config.Upstream{
					Address:        "127.0.0.1:1",
					PathRewrite:    "",
					ConnectTimeout: 100 * time.Millisecond,
					ReadTimeout:    100 * time.Millisecond,
					WriteTimeout:   100 * time.Millisecond,
					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()

	// Should get a 502 Bad Gateway when upstream is unreachable
	if resp.StatusCode != http.StatusBadGateway {
		t.Errorf("expected status 502, got %d", resp.StatusCode)
	}
}

func TestProxyMultipleRoutes(t *testing.T) {
	// Create two different upstream servers
	upstream1 := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusOK)
		fmt.Fprintf(w, "upstream1")
	}))
	defer upstream1.Close()

	upstream2 := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusOK)
		fmt.Fprintf(w, "upstream2")
	}))
	defer upstream2.Close()

	addr1 := strings.TrimPrefix(upstream1.URL, "http://")
	addr2 := strings.TrimPrefix(upstream2.URL, "http://")

	cfg := &config.Config{
		Routes: map[string]*config.Route{
			"route1": {
				Name: "route1",
				Upstream: config.Upstream{
					Address:        addr1,
					ConnectTimeout: 5 * time.Second,
					ReadTimeout:    5 * time.Second,
					WriteTimeout:   5 * time.Second,
					MaxBodySize:    1024 * 1024,
					AuthRequired:   false,
				},
			},
			"route2": {
				Name: "route2",
				Upstream: config.Upstream{
					Address:        addr2,
					ConnectTimeout: 5 * time.Second,
					ReadTimeout:    5 * time.Second,
					WriteTimeout:   5 * time.Second,
					MaxBodySize:    1024 * 1024,
					AuthRequired:   false,
				},
			},
		},
	}

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

	// With multiple routes, requests should be routed somewhere
	// (task 1.1 doesn't specify which route for undecorated requests,
	// but task 2.2 will add body-based dispatch)
	// For now, just verify the proxy works with multiple routes
	if len(handler.routes) != 2 {
		t.Errorf("expected 2 routes, got %d", len(handler.routes))
	}
}

func TestProxyPreservesMethod(t *testing.T) {
	method := ""
	upstreamServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		method = r.Method
		w.WriteHeader(http.StatusOK)
	}))
	defer upstreamServer.Close()

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

	cfg := &config.Config{
		Routes: map[string]*config.Route{
			"test-route": {
				Name: "test-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()

	methods := []string{"GET", "POST", "PUT", "DELETE"}
	for _, m := range methods {
		req, _ := http.NewRequest(m, server.URL+"/test", nil)
		resp, _ := http.DefaultClient.Do(req)
		resp.Body.Close()

		if method != m {
			t.Errorf("expected method %s, got %s", m, method)
		}
	}
}

func TestProxyPreservesQueryString(t *testing.T) {
	requestedURL := ""
	upstreamServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		requestedURL = r.URL.String()
		w.WriteHeader(http.StatusOK)
	}))
	defer upstreamServer.Close()

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

	cfg := &config.Config{
		Routes: map[string]*config.Route{
			"test-route": {
				Name: "test-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()

	http.Get(server.URL + "/test?key=value&other=param")

	if !strings.Contains(requestedURL, "key=value") || !strings.Contains(requestedURL, "other=param") {
		t.Errorf("query string not preserved: %s", requestedURL)
	}
}

func TestProxyPreservesBody(t *testing.T) {
	receivedBody := ""
	upstreamServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		body, _ := io.ReadAll(r.Body)
		receivedBody = string(body)
		w.WriteHeader(http.StatusOK)
	}))
	defer upstreamServer.Close()

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

	cfg := &config.Config{
		Routes: map[string]*config.Route{
			"test-route": {
				Name: "test-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()

	testBody := `{"model": "test", "messages": []}`
	resp, err := http.Post(server.URL+"/test", "application/json", strings.NewReader(testBody))
	if err != nil {
		t.Fatalf("request failed: %v", err)
	}
	resp.Body.Close()

	if receivedBody != testBody {
		t.Errorf("expected body %s, got %s", testBody, receivedBody)
	}
}

// TestProxyQueueFull verifies that a 503 is returned and the queue_full metric
// is incremented when the backend queue is at capacity.
//
// Setup: maxSlots=1, maxQueue=1. Request 1 holds the slot. Request 2 fills the
// queue (blocks waiting for slot). Request 3 gets ErrQueueFull → 503.
func TestProxyQueueFull(t *testing.T) {
	t.Setenv("VLLM_MAX_NUM_SEQS", "1")    // 1 concurrent slot
	t.Setenv("VLLM_MAX_QUEUE_DEPTH", "1") // queue capacity 1
	t.Setenv("WARMUP_DELAY_MS", "0")
	observability.Global().Reset()

	// Upstream blocks until released; buffers one signal when it first starts.
	slotHeld := make(chan struct{}, 1)
	release := make(chan struct{})
	upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		select {
		case slotHeld <- struct{}{}:
		default:
		}
		<-release
		w.WriteHeader(http.StatusOK)
	}))
	defer upstream.Close()

	upstreamAddr := strings.TrimPrefix(upstream.URL, "http://")
	cfg := &config.Config{
		Routes: make(map[string]*config.Route),
		Models: map[string]*config.ModelUpstream{
			"q-model": {
				Name:    "q-model",
				Address: upstreamAddr,
			},
		},
	}
	h := New(cfg)
	defer h.Close()

	srv := httptest.NewServer(h)
	defer srv.Close()

	mkReq := func() *http.Request {
		body := strings.NewReader(`{"model":"q-model","messages":[]}`)
		req, _ := http.NewRequest("POST", srv.URL+"/v1/chat/completions", body)
		req.Header.Set("Content-Type", "application/json")
		return req
	}

	// Request 1: acquires the only slot, reaches upstream, blocks there.
	r1ch := make(chan *http.Response, 1)
	go func() {
		resp, _ := http.DefaultClient.Do(mkReq())
		r1ch <- resp
	}()
	<-slotHeld // slot is acquired; upstream is holding request 1

	// Request 2: enters the queue (capacity 1), blocks waiting for slot.
	r2ch := make(chan *http.Response, 1)
	go func() {
		resp, _ := http.DefaultClient.Do(mkReq())
		r2ch <- resp
	}()
	// Spin until request 2 enters the backend queue channel.
	for len(h.schedulers["q-model"].backends[0].queue) == 0 {
		runtime.Gosched()
	}

	// Request 3: queue is full → expect 503 and metric increment.
	resp3, err := http.DefaultClient.Do(mkReq())
	if err != nil {
		t.Fatalf("request 3 failed: %v", err)
	}
	defer resp3.Body.Close()

	if resp3.StatusCode != http.StatusServiceUnavailable {
		t.Errorf("expected 503 queue full, got %d", resp3.StatusCode)
	}
	if got := observability.Global().GetQueueFull(); got != 1 {
		t.Errorf("expected GetQueueFull()=1, got %d", got)
	}

	// Cleanup: unblock upstream so goroutines can finish.
	close(release)
	if r := <-r1ch; r != nil {
		r.Body.Close()
	}
	if r := <-r2ch; r != nil {
		r.Body.Close()
	}
}

// TestProxyCacheHitLogging verifies that when a response contains the
// vllm-num-cached-tokens header and a cache key was set, the prefix cache
// hit token count is incremented in observability metrics.
func TestProxyCacheHitLogging(t *testing.T) {
	observability.Global().Reset()

	upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("vllm-num-cached-tokens", "128")
		w.WriteHeader(http.StatusOK)
	}))
	defer upstream.Close()

	upstreamAddr := strings.TrimPrefix(upstream.URL, "http://")
	cfg := &config.Config{
		Routes: make(map[string]*config.Route),
		Models: map[string]*config.ModelUpstream{
			"cache-model": {
				Name:    "cache-model",
				Address: upstreamAddr,
			},
		},
	}
	h := New(cfg)
	defer h.Close()

	srv := httptest.NewServer(h)
	defer srv.Close()

	body := strings.NewReader(`{"model":"cache-model","messages":[]}`)
	req, _ := http.NewRequest("POST", srv.URL+"/v1/chat/completions", body)
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("X-Poimen-Cache-Key", "testkey123")

	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		t.Fatalf("request failed: %v", err)
	}
	resp.Body.Close()

	if got := observability.Global().GetPrefixCacheHitTokens(); got != 128 {
		t.Errorf("expected PrefixCacheHitTokens=128, got %d", got)
	}
}

// TestWarmupHeaderHandling verifies that X-Poimen-Warmup: true is stripped from
// the upstream request and replaced with X-Priority: low, while a request without
// X-Poimen-Warmup reaches upstream without either header.
//
// The warmup path runs only for model-based routes (cfg.Models + scheduler),
// so we POST to /v1/chat/completions with a JSON body that names the test model.
// WARMUP_DELAY_MS=0 eliminates the 500 ms sleep.
func TestWarmupHeaderHandling(t *testing.T) {
	t.Setenv("WARMUP_DELAY_MS", "0")

	tests := []struct {
		name            string
		sendWarmup      bool
		wantWarmupGone  bool
		wantPriorityLow bool
	}{
		{
			name:            "warmup=true",
			sendWarmup:      true,
			wantWarmupGone:  true, // X-Poimen-Warmup must NOT reach upstream
			wantPriorityLow: true, // X-Priority: low must be set
		},
		{
			name:            "warmup=false",
			sendWarmup:      false,
			wantWarmupGone:  true,  // no header sent, so absent at upstream
			wantPriorityLow: false, // X-Priority must NOT be set
		},
	}

	for _, tc := range tests {
		t.Run(tc.name, func(t *testing.T) {
			var gotWarmup, gotPriority string
			upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
				gotWarmup = r.Header.Get("X-Poimen-Warmup")
				gotPriority = r.Header.Get("X-Priority")
				w.WriteHeader(http.StatusOK)
			}))
			defer upstream.Close()

			upstreamAddr := strings.TrimPrefix(upstream.URL, "http://")
			// Warmup code only runs for model-based routes (scheduler path).
			// Use cfg.Models so New() builds a Scheduler for "warmup-model".
			cfg := &config.Config{
				Routes: make(map[string]*config.Route),
				Models: map[string]*config.ModelUpstream{
					"warmup-model": {
						Name:    "warmup-model",
						Address: upstreamAddr,
					},
				},
			}
			h := New(cfg)
			defer h.Close()

			srv := httptest.NewServer(h)
			defer srv.Close()

			// POST to /v1/chat/completions — triggers model-based dispatch and warmup path.
			body := strings.NewReader(`{"model":"warmup-model","messages":[]}`)
			req, _ := http.NewRequest("POST", srv.URL+"/v1/chat/completions", body)
			req.Header.Set("Content-Type", "application/json")
			if tc.sendWarmup {
				req.Header.Set("X-Poimen-Warmup", "true")
			}
			resp, err := http.DefaultClient.Do(req)
			if err != nil {
				t.Fatalf("request failed: %v", err)
			}
			resp.Body.Close()

			if tc.wantWarmupGone && gotWarmup != "" {
				t.Errorf("X-Poimen-Warmup must be stripped before upstream, got %q", gotWarmup)
			}
			if tc.wantPriorityLow && gotPriority != "low" {
				t.Errorf("X-Priority: low must be set when warmup=true, got %q", gotPriority)
			}
			if !tc.wantPriorityLow && gotPriority != "" {
				t.Errorf("X-Priority must not be set for normal requests, got %q", gotPriority)
			}
		})
	}
}