package proxy import ( "fmt" "io" "net/http" "net/http/httptest" "strings" "testing" "time" "forgejo.riotpiao.com/rock/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() // 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) } }