Improvements for the vmcp e2e test infrastructure (#5026)

test/e2e: fix polling loop cancellation, unify OIDC options, remove stale link

- Replace time.Now()/time.Sleep polling in WaitForVMCPHealthReady with
  context.WithTimeout + time.NewTicker + select, matching the pattern
  already used in WaitForMCPServerReady
- Unify NewOIDCMockServer and NewOIDCMockServerWithClientOptions into a
  single variadic constructor accepting OIDCMockOption, which covers both
  Fosite-config and client-registration axes; update WithClientAudience
  and WithAccessTokenLifespan to return OIDCMockOption
- Remove the Tracked by issue link from the vmcp_infra_features_test.go
  package comment (belongs in the PR/commit, not in source)

Closes #5022

Co-authored-by: taskbot <taskbot@users.noreply.github.com>
Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com>

Author: 3 people

test/e2e/mcp_client_helpers.go

@@ -69,28 +69,39 @@ func NewMCPClientForStreamableHTTPWithToken(config *TestConfig, serverURL, token
 // is configured (MCP Initialize would fail with 401 for unauthenticated probes).
 func WaitForVMCPHealthReady(healthURL string, timeout time.Duration) error {
 	httpClient := &http.Client{Timeout: 5 * time.Second}
-	deadline := time.Now().Add(timeout)
+	ctx, cancel := context.WithTimeout(context.Background(), timeout)
+	defer cancel()
+	ticker := time.NewTicker(2 * time.Second)
+	defer ticker.Stop()
 	var lastErr error
 	var lastStatus int
 	var lastBody string
-	for time.Now().Before(deadline) {
-		resp, err := httpClient.Get(healthURL) //nolint:gosec // URL is test-controlled
-		if resp != nil {
-			lastStatus = resp.StatusCode
-			bodyBytes, _ := io.ReadAll(io.LimitReader(resp.Body, 512))
-			_ = resp.Body.Close()
-			lastBody = string(bodyBytes)
-		}
-		lastErr = err
-		if err == nil && resp.StatusCode == http.StatusOK {
-			return nil
+	for {
+		select {
+		case <-ctx.Done():
+			if lastErr != nil {
+				return fmt.Errorf("timeout waiting for vMCP health endpoint at %s: last error: %w", healthURL, lastErr)
+			}
+			return fmt.Errorf("timeout waiting for vMCP health endpoint at %s: last status: %d, body: %s", healthURL, lastStatus, lastBody)
+		case <-ticker.C:
+			req, err := http.NewRequestWithContext(ctx, http.MethodGet, healthURL, nil) //nolint:gosec // URL is test-controlled
+			if err != nil {
+				lastErr = err
+				continue
+			}
+			resp, err := httpClient.Do(req)
+			if resp != nil {
+				lastStatus = resp.StatusCode
+				bodyBytes, _ := io.ReadAll(io.LimitReader(resp.Body, 512))
+				_ = resp.Body.Close()
+				lastBody = string(bodyBytes)
+			}
+			lastErr = err
+			if err == nil && resp.StatusCode == http.StatusOK {
+				return nil
+			}
 		}
-		time.Sleep(2 * time.Second)
-	}
-	if lastErr != nil {
-		return fmt.Errorf("timeout waiting for vMCP health endpoint at %s: last error: %w", healthURL, lastErr)
 	}
-	return fmt.Errorf("timeout waiting for vMCP health endpoint at %s: last status: %d, body: %s", healthURL, lastStatus, lastBody)
 }
 
 // Initialize initializes the MCP connection

test/e2e/oidc_mock.go

@@ -50,33 +50,35 @@ type AuthRequest struct {
 // pkg/auth's token validator requires the kid claim to look up the signing key.
 const jwtKeyID = "test-key-1"
 
-// OIDCMockServerOption configures the OIDCMockServer client registration.
-type OIDCMockServerOption func(*fosite.DefaultClient)
+// OIDCMockOption is a unified option type for configuring the OIDC mock server.
+// Use WithClientAudience for client-registration settings and WithAccessTokenLifespan
+// (or other fosite-level helpers) for token-lifecycle settings. A single constructor
+// accepts both, so tests needing both a custom token lifetime and a specific audience
+// no longer require separate constructors.
+type OIDCMockOption struct {
+	fositeOpt func(*fosite.Config)
+	clientOpt func(*fosite.DefaultClient)
+}
 
 // WithClientAudience sets the allowed audience(s) on the registered test client.
 // Use this when the vMCP OIDC config requires a specific audience claim in tokens.
-func WithClientAudience(audiences ...string) OIDCMockServerOption {
-	return func(c *fosite.DefaultClient) {
+func WithClientAudience(audiences ...string) OIDCMockOption {
+	return OIDCMockOption{clientOpt: func(c *fosite.DefaultClient) {
 		c.Audience = audiences
-	}
+	}}
 }
 
-// NewOIDCMockServer creates a new OIDC mock server using Ory Fosite
-func NewOIDCMockServer(port int, clientID, clientSecret string, opts ...func(*fosite.Config)) (*OIDCMockServer, error) {
+// NewOIDCMockServer creates a new OIDC mock server using Ory Fosite.
+// Use WithClientAudience to set client-level options and WithAccessTokenLifespan
+// for Fosite-level settings. Both option kinds may be mixed in a single call.
+func NewOIDCMockServer(port int, clientID, clientSecret string, opts ...OIDCMockOption) (*OIDCMockServer, error) {
 	config := defaultFositeConfig(port)
 	for _, opt := range opts {
-		opt(config)
+		if opt.fositeOpt != nil {
+			opt.fositeOpt(config)
+		}
 	}
-	return newOIDCMockServer(port, clientID, clientSecret, config)
-}
-
-// NewOIDCMockServerWithClientOptions creates a new OIDC mock server, applying
-// clientOpts to the registered test client. Use this when you need to control
-// client-level settings such as Audience.
-func NewOIDCMockServerWithClientOptions(
-	port int, clientID, clientSecret string, clientOpts ...OIDCMockServerOption,
-) (*OIDCMockServer, error) {
-	return newOIDCMockServer(port, clientID, clientSecret, defaultFositeConfig(port), clientOpts...)
+	return newOIDCMockServer(port, clientID, clientSecret, config, opts...)
 }
 
 // defaultFositeConfig returns the standard Fosite config for the mock server.
@@ -93,9 +95,9 @@ func defaultFositeConfig(port int) *fosite.Config {
 	}
 }
 
-// newOIDCMockServer is the shared implementation for the public constructors.
+// newOIDCMockServer is the shared implementation for NewOIDCMockServer.
 func newOIDCMockServer(
-	port int, clientID, clientSecret string, config *fosite.Config, clientOpts ...OIDCMockServerOption,
+	port int, clientID, clientSecret string, config *fosite.Config, opts ...OIDCMockOption,
 ) (*OIDCMockServer, error) {
 	// Generate RSA key for JWT signing
 	key, err := rsa.GenerateKey(rand.Reader, 2048)
@@ -121,8 +123,10 @@ func newOIDCMockServer(
 		GrantTypes:    []string{"authorization_code", "refresh_token", "client_credentials"},
 		Scopes:        []string{"openid", "profile", "email"},
 	}
-	for _, opt := range clientOpts {
-		opt(client)
+	for _, opt := range opts {
+		if opt.clientOpt != nil {
+			opt.clientOpt(client)
+		}
 	}
 	store.Clients[clientID] = client
 
@@ -497,8 +501,8 @@ func (*OIDCMockServer) CompleteAuthRequest(authReq *AuthRequest) error {
 }
 
 // WithAccessTokenLifespan sets the lifespan of access tokens for the OIDC mock server.
-func WithAccessTokenLifespan(d time.Duration) func(*fosite.Config) {
-	return func(c *fosite.Config) {
+func WithAccessTokenLifespan(d time.Duration) OIDCMockOption {
+	return OIDCMockOption{fositeOpt: func(c *fosite.Config) {
 		c.AccessTokenLifespan = d
-	}
+	}}
 }

test/e2e/vmcp_infra_features_test.go

@@ -4,7 +4,6 @@
 // Package e2e_test contains infrastructure-heavy vMCP CLI e2e tests that require
 // external services (OIDC server, Redis) as test fixtures.
 // These complement the basic feature tests in vmcp_cli_features_test.go.
-// Tracked by: https://github.com/stacklok/toolhive/issues/4944
 package e2e_test
 
 import (
@@ -107,7 +106,7 @@ var _ = Describe("vMCP infra features", Label("vmcp", "e2e", "infra"), func() {
 
 			oidcPort = allocateVMCPPort()
 			var err error
-			oidcServer, err = e2e.NewOIDCMockServerWithClientOptions(oidcPort, "test-client", "test-secret",
+			oidcServer, err = e2e.NewOIDCMockServer(oidcPort, "test-client", "test-secret",
 				e2e.WithClientAudience("vmcp-e2e-test"),
 			)
 			Expect(err).ToNot(HaveOccurred())