homelab-frontend/internal/proxy/proxy_test.go

package proxy

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

	"github.com/Riotpiaole/homelab-frontend/internal/config"
)

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()

	resp, err := http.Get(server.URL + "/v1/models")
	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, _ := http.Post(server.URL+"/test", "application/json", strings.NewReader(testBody))
	resp.Body.Close()

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