refactor: add aiproxy test helpers to reduce boilerplate (#21404)

## Description Adds test helper functions to reduce boilerplate in `aibridgeproxyd` tests: * `newTestProxy`: creates a proxy server with functional options, waits for it to be ready * `newProxyClient`: creates an HTTP client configured to use the proxy * `newTargetServer`: creates a mock HTTPS server and returns its URL Related to: https://github.com/coder/coder/pull/21344#discussion_r2638930199
2026-09-21 20:51:01 +08:00 · 2025-12-31 10:15:59 +00:00
parent 9ac865b72f
commit e5377fbd93
1 changed files with 174 additions and 258 deletions
@@ -96,6 +96,131 @@ func generateSharedTestCA() (certFile, keyFile string, err error) {
 	return certPath, keyPath, nil
 }

+type testProxyConfig struct {
+	listenAddr     string
+	coderAccessURL string
+	allowedPorts   []string
+	certStore      *aibridgeproxyd.CertCache
+}
+
+type testProxyOption func(*testProxyConfig)
+
+func withAllowedPorts(ports ...string) testProxyOption {
+	return func(cfg *testProxyConfig) {
+		cfg.allowedPorts = ports
+	}
+}
+
+func withCoderAccessURL(coderAccessURL string) testProxyOption {
+	return func(cfg *testProxyConfig) {
+		cfg.coderAccessURL = coderAccessURL
+	}
+}
+
+func withCertStore(store *aibridgeproxyd.CertCache) testProxyOption {
+	return func(cfg *testProxyConfig) {
+		cfg.certStore = store
+	}
+}
+
+// newTestProxy creates a new AI Bridge Proxy server for testing.
+// It uses the shared test CA and registers cleanup automatically.
+// It waits for the proxy server to be ready before returning.
+func newTestProxy(t *testing.T, opts ...testProxyOption) *aibridgeproxyd.Server {
+	t.Helper()
+
+	cfg := &testProxyConfig{
+		listenAddr:     "127.0.0.1:0",
+		coderAccessURL: "http://localhost:3000",
+	}
+	for _, opt := range opts {
+		opt(cfg)
+	}
+
+	certFile, keyFile := getSharedTestCA(t)
+	logger := slogtest.Make(t, nil)
+
+	aibridgeOpts := aibridgeproxyd.Options{
+		ListenAddr:     cfg.listenAddr,
+		CoderAccessURL: cfg.coderAccessURL,
+		CertFile:       certFile,
+		KeyFile:        keyFile,
+		AllowedPorts:   cfg.allowedPorts,
+	}
+	if cfg.certStore != nil {
+		aibridgeOpts.CertStore = cfg.certStore
+	}
+
+	srv, err := aibridgeproxyd.New(t.Context(), logger, aibridgeOpts)
+	require.NoError(t, err)
+	t.Cleanup(func() { _ = srv.Close() })
+
+	// Wait for the proxy server to be ready.
+	proxyAddr := srv.Addr()
+	require.NotEmpty(t, proxyAddr)
+	require.Eventually(t, func() bool {
+		conn, err := net.Dial("tcp", proxyAddr)
+		if err != nil {
+			return false
+		}
+		_ = conn.Close()
+		return true
+	}, testutil.WaitShort, testutil.IntervalFast)
+
+	return srv
+}
+
+// newProxyClient creates an HTTP client configured to use the proxy and trust its CA.
+// It adds a Proxy-Authorization header with the provided token for authentication.
+func newProxyClient(t *testing.T, srv *aibridgeproxyd.Server, proxyAuth string) *http.Client {
+	t.Helper()
+
+	certFile, _ := getSharedTestCA(t)
+
+	// Load the CA certificate so the client trusts the proxy's MITM certificate.
+	certPEM, err := os.ReadFile(certFile)
+	require.NoError(t, err)
+	certPool := x509.NewCertPool()
+	ok := certPool.AppendCertsFromPEM(certPEM)
+	require.True(t, ok)
+
+	// Create an HTTP client configured to use the proxy.
+	proxyURL, err := url.Parse("http://" + srv.Addr())
+	require.NoError(t, err)
+
+	transport := &http.Transport{
+		Proxy: http.ProxyURL(proxyURL),
+		TLSClientConfig: &tls.Config{
+			MinVersion: tls.VersionTLS12,
+			RootCAs:    certPool,
+		},
+	}
+
+	// Only set the header if proxyAuth is provided. This allows tests to
+	// verify behavior when the Proxy-Authorization header is missing.
+	if proxyAuth != "" {
+		transport.ProxyConnectHeader = http.Header{
+			"Proxy-Authorization": []string{proxyAuth},
+		}
+	}
+
+	return &http.Client{Transport: transport}
+}
+
+// newTargetServer creates a mock HTTPS server that will be the target of proxied requests.
+// It returns the server's parsed URL. The server is automatically closed when the test ends.
+func newTargetServer(t *testing.T, handler http.HandlerFunc) *url.URL {
+	t.Helper()
+
+	srv := httptest.NewTLSServer(handler)
+	t.Cleanup(srv.Close)
+
+	srvURL, err := url.Parse(srv.URL)
+	require.NoError(t, err)
+
+	return srvURL
+}
+
 // makeProxyAuthHeader creates a Proxy-Authorization header value with the given token.
 // Format: "Basic base64(username:token)" where username is "ignored".
 func makeProxyAuthHeader(token string) string {
@@ -257,72 +382,24 @@ func TestClose(t *testing.T) {
 func TestProxy_CertCaching(t *testing.T) {
 	t.Parallel()

-	// Create a mock HTTPS server that will be the target of our proxied request.
-	targetServer := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+	// Create a mock HTTPS server that will be the target of the proxied request.
+	targetURL := newTargetServer(t, func(w http.ResponseWriter, r *http.Request) {
 		w.WriteHeader(http.StatusOK)
-	}))
-	defer targetServer.Close()
-
-	targetURL, err := url.Parse(targetServer.URL)
-	require.NoError(t, err)
-
-	certFile, keyFile := getSharedTestCA(t)
-	logger := slogtest.Make(t, nil)
+	})

 	// Create a cert cache so we can inspect it after the request.
 	certCache := aibridgeproxyd.NewCertCache()

 	// Start the proxy server with the certificate cache.
-	srv, err := aibridgeproxyd.New(t.Context(), logger, aibridgeproxyd.Options{
-		ListenAddr:     "127.0.0.1:0",
-		CoderAccessURL: "http://localhost:3000",
-		CertFile:       certFile,
-		KeyFile:        keyFile,
-		CertStore:      certCache,
-		AllowedPorts:   []string{targetURL.Port()},
-	})
-	require.NoError(t, err)
-	t.Cleanup(func() { _ = srv.Close() })
-
-	proxyAddr := srv.Addr()
-	require.NotEmpty(t, proxyAddr)
-
-	// Wait for the proxy server to be ready.
-	require.Eventually(t, func() bool {
-		conn, err := net.Dial("tcp", proxyAddr)
-		if err != nil {
-			return false
-		}
-		_ = conn.Close()
-		return true
-	}, testutil.WaitShort, testutil.IntervalFast)
-
-	// Load the CA certificate so the client trusts the proxy's MITM certificate.
-	certPEM, err := os.ReadFile(certFile)
-	require.NoError(t, err)
-	certPool := x509.NewCertPool()
-	certPool.AppendCertsFromPEM(certPEM)
-
-	// Create an HTTP client configured to use the proxy.
-	proxyURL, err := url.Parse("http://" + proxyAddr)
-	require.NoError(t, err)
-
-	client := &http.Client{
-		Transport: &http.Transport{
-			Proxy: http.ProxyURL(proxyURL),
-			ProxyConnectHeader: http.Header{
-				"Proxy-Authorization": []string{makeProxyAuthHeader("test-session-token")},
-			},
-			TLSClientConfig: &tls.Config{
-				MinVersion: tls.VersionTLS12,
-				RootCAs:    certPool,
-			},
-		},
-	}
+	srv := newTestProxy(t,
+		withAllowedPorts(targetURL.Port()),
+		withCertStore(certCache),
+	)

 	// Make a request through the proxy to the target server.
 	// This triggers MITM and caches the generated certificate.
-	req, err := http.NewRequestWithContext(t.Context(), http.MethodGet, targetServer.URL, nil)
+	client := newProxyClient(t, srv, makeProxyAuthHeader("test-session-token"))
+	req, err := http.NewRequestWithContext(t.Context(), http.MethodGet, targetURL.String(), nil)
 	require.NoError(t, err)
 	resp, err := client.Do(req)
 	require.NoError(t, err)
@@ -343,17 +420,23 @@ func TestProxy_PortValidation(t *testing.T) {
 	t.Parallel()

 	tests := []struct {
-		name        string
-		allowPort   bool
-		expectError bool
+		name         string
+		allowedPorts func(targetURL *url.URL) []string
+		expectError  bool
 	}{
 		{
-			name:      "AllowedPort",
-			allowPort: true,
+			name: "AllowedPort",
+			// Include the target's random port so the request is allowed.
+			allowedPorts: func(targetURL *url.URL) []string {
+				return []string{targetURL.Port()}
+			},
 		},
 		{
-			name:        "RejectedPort",
-			allowPort:   false,
+			name: "RejectedPort",
+			// Only allow port 443 which doesn't match the target.
+			allowedPorts: func(_ *url.URL) []string {
+				return []string{"443"}
+			},
 			expectError: true,
 		},
 	}
@@ -363,77 +446,17 @@ func TestProxy_PortValidation(t *testing.T) {
 			t.Parallel()

 			// Create a target HTTPS server that will be the destination of our proxied request.
-			targetServer := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			targetURL := newTargetServer(t, func(w http.ResponseWriter, r *http.Request) {
 				w.WriteHeader(http.StatusOK)
 				_, _ = w.Write([]byte("hello from target"))
-			}))
-			t.Cleanup(func() { targetServer.Close() })
-
-			targetURL, err := url.Parse(targetServer.URL)
-			require.NoError(t, err)
-
-			certFile, keyFile := getSharedTestCA(t)
-			logger := slogtest.Make(t, nil)
-
-			// Configure allowed ports based on test case.
-			// For allowed case, include the target's random port.
-			// For rejected case, only allow port 443 which doesn't match the target.
-			var allowedPorts []string
-			if tt.allowPort {
-				allowedPorts = []string{targetURL.Port()}
-			} else {
-				allowedPorts = []string{"443"}
-			}
+			})

 			// Start the proxy server on a random port to avoid conflicts when running tests in parallel.
-			srv, err := aibridgeproxyd.New(t.Context(), logger, aibridgeproxyd.Options{
-				ListenAddr:     "127.0.0.1:0",
-				CoderAccessURL: "http://localhost:3000",
-				CertFile:       certFile,
-				KeyFile:        keyFile,
-				AllowedPorts:   allowedPorts,
-			})
-			require.NoError(t, err)
-			t.Cleanup(func() { _ = srv.Close() })
-
-			proxyAddr := srv.Addr()
-			require.NotEmpty(t, proxyAddr)
-
-			// Wait for the proxy server to be ready.
-			require.Eventually(t, func() bool {
-				conn, err := net.Dial("tcp", proxyAddr)
-				if err != nil {
-					return false
-				}
-				_ = conn.Close()
-				return true
-			}, testutil.WaitShort, testutil.IntervalFast)
-
-			// Load the CA certificate so the client trusts the proxy's MITM certificate.
-			certPEM, err := os.ReadFile(certFile)
-			require.NoError(t, err)
-			certPool := x509.NewCertPool()
-			certPool.AppendCertsFromPEM(certPEM)
-
-			// Create an HTTP client configured to use the proxy.
-			proxyURL, err := url.Parse("http://" + proxyAddr)
-			require.NoError(t, err)
-
-			client := &http.Client{
-				Transport: &http.Transport{
-					Proxy: http.ProxyURL(proxyURL),
-					ProxyConnectHeader: http.Header{
-						"Proxy-Authorization": []string{makeProxyAuthHeader("test-session-token")},
-					},
-					TLSClientConfig: &tls.Config{
-						MinVersion: tls.VersionTLS12,
-						RootCAs:    certPool,
-					},
-				},
-			}
+			srv := newTestProxy(t, withAllowedPorts(tt.allowedPorts(targetURL)...))

 			// Make a request through the proxy to the target server.
-			req, err := http.NewRequestWithContext(t.Context(), http.MethodGet, targetServer.URL, nil)
+			client := newProxyClient(t, srv, makeProxyAuthHeader("test-session-token"))
+			req, err := http.NewRequestWithContext(t.Context(), http.MethodGet, targetURL.String(), nil)
 			require.NoError(t, err)

 			resp, err := client.Do(req)
@@ -445,6 +468,7 @@ func TestProxy_PortValidation(t *testing.T) {
 			require.NoError(t, err)
 			defer resp.Body.Close()

+			// Verify the request was successful and reached the target server.
 			body, err := io.ReadAll(resp.Body)
 			require.NoError(t, err)
 			require.Equal(t, http.StatusOK, resp.StatusCode)
@@ -488,71 +512,17 @@ func TestProxy_Authentication(t *testing.T) {
 			t.Parallel()

 			// Create a mock HTTPS server that will be the target of our proxied request.
-			targetServer := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			targetURL := newTargetServer(t, func(w http.ResponseWriter, _ *http.Request) {
 				w.WriteHeader(http.StatusOK)
 				_, _ = w.Write([]byte("hello from target"))
-			}))
-			t.Cleanup(func() { targetServer.Close() })
-
-			targetURL, err := url.Parse(targetServer.URL)
-			require.NoError(t, err)
-
-			certFile, keyFile := getSharedTestCA(t)
-			logger := slogtest.Make(t, nil)
+			})

 			// Start the proxy server on a random port to avoid conflicts when running tests in parallel.
-			// The actual port is accessible via srv.Addr().
-			srv, err := aibridgeproxyd.New(t.Context(), logger, aibridgeproxyd.Options{
-				ListenAddr:     "127.0.0.1:0",
-				CoderAccessURL: "http://localhost:3000",
-				CertFile:       certFile,
-				KeyFile:        keyFile,
-				AllowedPorts:   []string{targetURL.Port()},
-			})
-			require.NoError(t, err)
-			t.Cleanup(func() { _ = srv.Close() })
-
-			proxyAddr := srv.Addr()
-			require.NotEmpty(t, proxyAddr)
-
-			// Wait for the proxy server to be ready.
-			require.Eventually(t, func() bool {
-				conn, err := net.Dial("tcp", proxyAddr)
-				if err != nil {
-					return false
-				}
-				_ = conn.Close()
-				return true
-			}, testutil.WaitShort, testutil.IntervalFast)
-
-			// Load the CA certificate so the client trusts the proxy's MITM certificate.
-			certPEM, err := os.ReadFile(certFile)
-			require.NoError(t, err)
-			certPool := x509.NewCertPool()
-			certPool.AppendCertsFromPEM(certPEM)
-
-			// Create an HTTP client configured to use the proxy.
-			proxyURL, err := url.Parse("http://" + proxyAddr)
-			require.NoError(t, err)
-
-			transport := &http.Transport{
-				Proxy: http.ProxyURL(proxyURL),
-				TLSClientConfig: &tls.Config{
-					MinVersion: tls.VersionTLS12,
-					RootCAs:    certPool,
-				},
-			}
-
-			if tt.proxyAuth != "" {
-				transport.ProxyConnectHeader = http.Header{
-					"Proxy-Authorization": []string{tt.proxyAuth},
-				}
-			}
-
-			client := &http.Client{Transport: transport}
+			srv := newTestProxy(t, withAllowedPorts(targetURL.Port()))

 			// Make a request through the proxy to the target server.
-			req, err := http.NewRequestWithContext(t.Context(), http.MethodGet, targetServer.URL, nil)
+			client := newProxyClient(t, srv, tt.proxyAuth)
+			req, err := http.NewRequestWithContext(t.Context(), http.MethodGet, targetURL.String(), nil)
 			require.NoError(t, err)
 			resp, err := client.Do(req)

@@ -624,7 +594,7 @@ func TestProxy_MITM(t *testing.T) {
 			var receivedPath string
 			var receivedAuth string

-			// Create a mock aibridged server.
+			// Create a mock aibridged server that captures requests.
 			aibridgedServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 				receivedPath = r.URL.Path
 				receivedAuth = r.Header.Get("Authorization")
@@ -634,14 +604,10 @@ func TestProxy_MITM(t *testing.T) {
 			t.Cleanup(func() { aibridgedServer.Close() })

 			// Create a mock target server for passthrough tests.
-			targetServer := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			passthroughURL := newTargetServer(t, func(w http.ResponseWriter, _ *http.Request) {
 				w.WriteHeader(http.StatusOK)
 				_, _ = w.Write([]byte("hello from passthrough"))
-			}))
-			t.Cleanup(func() { targetServer.Close() })
-
-			certFile, keyFile := getSharedTestCA(t)
-			logger := slogtest.Make(t, nil)
+			})

 			// Configure allowed ports based on test case.
 			// AI provider tests connect to the specified port, or 443 if not specified.
@@ -649,9 +615,7 @@ func TestProxy_MITM(t *testing.T) {
 			var allowedPorts []string
 			switch {
 			case tt.passthrough:
-				parsedTargetURL, err := url.Parse(targetServer.URL)
-				require.NoError(t, err)
-				allowedPorts = []string{parsedTargetURL.Port()}
+				allowedPorts = []string{passthroughURL.Port()}
 			case tt.targetPort != "":
 				allowedPorts = []string{tt.targetPort}
 			default:
@@ -659,60 +623,20 @@ func TestProxy_MITM(t *testing.T) {
 			}

 			// Start the proxy server pointing to our mock aibridged.
-			srv, err := aibridgeproxyd.New(t.Context(), logger, aibridgeproxyd.Options{
-				ListenAddr:     "127.0.0.1:0",
-				CoderAccessURL: aibridgedServer.URL,
-				CertFile:       certFile,
-				KeyFile:        keyFile,
-				AllowedPorts:   allowedPorts,
-			})
-			require.NoError(t, err)
-			t.Cleanup(func() { _ = srv.Close() })
-
-			proxyAddr := srv.Addr()
-			require.NotEmpty(t, proxyAddr)
-
-			// Wait for the proxy server to be ready.
-			require.Eventually(t, func() bool {
-				conn, err := net.Dial("tcp", proxyAddr)
-				if err != nil {
-					return false
-				}
-				_ = conn.Close()
-				return true
-			}, testutil.WaitShort, testutil.IntervalFast)
-
-			// Load the CA certificate.
-			certPEM, err := os.ReadFile(certFile)
-			require.NoError(t, err)
-			certPool := x509.NewCertPool()
-			certPool.AppendCertsFromPEM(certPEM)
-
-			// Create an HTTP client configured to use the proxy.
-			proxyURL, err := url.Parse("http://" + proxyAddr)
-			require.NoError(t, err)
-
-			client := &http.Client{
-				Transport: &http.Transport{
-					Proxy: http.ProxyURL(proxyURL),
-					ProxyConnectHeader: http.Header{
-						"Proxy-Authorization": []string{makeProxyAuthHeader("test-session-token")},
-					},
-					TLSClientConfig: &tls.Config{
-						MinVersion: tls.VersionTLS12,
-						RootCAs:    certPool,
-					},
-				},
-			}
+			srv := newTestProxy(t,
+				withCoderAccessURL(aibridgedServer.URL),
+				withAllowedPorts(allowedPorts...),
+			)

 			// Build the target URL:
-			// - For passthrough, target the local mock TLS server.
-			// - For AI providers, use their real hostnames to trigger routing.
-			//   Non-default ports are included explicitly; default port (443) is omitted.
+			//   - For passthrough, target the local mock TLS server.
+			//   - For AI providers, use their real hostnames to trigger routing.
+			//     Non-default ports are included explicitly; default port (443) is omitted.
 			var targetURL string
+			var err error
 			switch {
 			case tt.passthrough:
-				targetURL, err = url.JoinPath(targetServer.URL, tt.targetPath)
+				targetURL, err = url.JoinPath(passthroughURL.String(), tt.targetPath)
 				require.NoError(t, err)
 			case tt.targetPort != "":
 				targetURL, err = url.JoinPath("https://"+tt.targetHost+":"+tt.targetPort, tt.targetPath)
@@ -723,6 +647,7 @@ func TestProxy_MITM(t *testing.T) {
 			}

 			// Make a request through the proxy to the target URL.
+			client := newProxyClient(t, srv, makeProxyAuthHeader("test-session-token"))
 			req, err := http.NewRequestWithContext(t.Context(), http.MethodPost, targetURL, strings.NewReader(`{}`))
 			require.NoError(t, err)
 			req.Header.Set("Content-Type", "application/json")
@@ -761,17 +686,7 @@ func TestServeCACert(t *testing.T) {
 	t.Run("Success", func(t *testing.T) {
 		t.Parallel()

-		certFile, keyFile := getSharedTestCA(t)
-		logger := slogtest.Make(t, nil)
-
-		srv, err := aibridgeproxyd.New(t.Context(), logger, aibridgeproxyd.Options{
-			ListenAddr:     "127.0.0.1:0",
-			CoderAccessURL: "http://localhost:3000",
-			CertFile:       certFile,
-			KeyFile:        keyFile,
-		})
-		require.NoError(t, err)
-		t.Cleanup(func() { _ = srv.Close() })
+		srv := newTestProxy(t)

 		// Create a request to the CA cert endpoint via the Handler.
 		req := httptest.NewRequest(http.MethodGet, "/ca-cert.pem", nil)
@@ -795,6 +710,7 @@ func TestServeCACert(t *testing.T) {
 		require.NotNil(t, cert)

 		// Verify it matches the original certificate.
+		certFile, _ := getSharedTestCA(t)
 		expectedCertPEM, err := os.ReadFile(certFile)
 		require.NoError(t, err)
 		require.Equal(t, expectedCertPEM, body)