extauth: refactor enterprise auth tests (#879)

auth/extauth/enterprise_handler_test.go

@@ -1,43 +1,37 @@
 // Copyright 2026 The Go MCP SDK Authors. All rights reserved.
-// Use of this source code is governed by an MIT-style
-// license that can be found in the LICENSE file.
+// Use of this source code is governed by the license
+// that can be found in the LICENSE file.
 
 package extauth
 
 import (
 	"context"
-	"encoding/json"
 	"fmt"
 	"net/http"
 	"net/http/httptest"
 	"strings"
 	"testing"
 
+	"github.com/modelcontextprotocol/go-sdk/internal/oauthtest"
 	"github.com/modelcontextprotocol/go-sdk/oauthex"
 	"golang.org/x/oauth2"
 )
 
-// TestNewEnterpriseHandler_Validation tests validation in NewEnterpriseHandler.
-func TestNewEnterpriseHandler_Validation(t *testing.T) {
-	validConfig := &EnterpriseHandlerConfig{
-		IdPIssuerURL: "https://idp.example.com",
-		IdPCredentials: &oauthex.ClientCredentials{
-			ClientID: "idp_client_id",
-		},
+func validEnterpriseHandlerConfig() *EnterpriseHandlerConfig {
+	return &EnterpriseHandlerConfig{
+		IdPIssuerURL:     "https://idp.example.com",
+		IdPCredentials:   &oauthex.ClientCredentials{ClientID: "idp_client_id"},
 		MCPAuthServerURL: "https://mcp-auth.example.com",
 		MCPResourceURI:   "https://mcp.example.com",
-		MCPCredentials: &oauthex.ClientCredentials{
-			ClientID: "mcp_client_id",
-		},
+		MCPCredentials:   &oauthex.ClientCredentials{ClientID: "mcp_client_id"},
 		IDTokenFetcher: func(ctx context.Context) (*oauth2.Token, error) {
-			token := &oauth2.Token{
-				AccessToken: "mock_access_token",
-				TokenType:   "Bearer",
-			}
+			token := &oauth2.Token{AccessToken: "mock", TokenType: "Bearer"}
 			return token.WithExtra(map[string]any{"id_token": "mock_id_token"}), nil
 		},
 	}
+}
 
+func TestNewEnterpriseHandler_Validation(t *testing.T) {
 	tests := []struct {
 		name      string
 		config    *EnterpriseHandlerConfig
@@ -50,151 +44,101 @@ func TestNewEnterpriseHandler_Validation(t *testing.T) {
 		},
 		{
 			name: "missing IdPIssuerURL",
-			config: &EnterpriseHandlerConfig{
-				IdPCredentials:   &oauthex.ClientCredentials{ClientID: "id"},
-				MCPAuthServerURL: "https://mcp-auth.example.com",
-				MCPResourceURI:   "https://mcp.example.com",
-				MCPCredentials:   &oauthex.ClientCredentials{ClientID: "id"},
-				IDTokenFetcher:   func(ctx context.Context) (*oauth2.Token, error) { return nil, nil },
-			},
+			config: func() *EnterpriseHandlerConfig {
+				c := validEnterpriseHandlerConfig()
+				c.IdPIssuerURL = ""
+				return c
+			}(),
 			wantError: "IdPIssuerURL is required",
 		},
 		{
 			name: "nil IdPCredentials",
-			config: &EnterpriseHandlerConfig{
-				IdPIssuerURL:     "https://idp.example.com",
-				IdPCredentials:   nil,
-				MCPAuthServerURL: "https://mcp-auth.example.com",
-				MCPResourceURI:   "https://mcp.example.com",
-				MCPCredentials:   &oauthex.ClientCredentials{ClientID: "id"},
-				IDTokenFetcher:   func(ctx context.Context) (*oauth2.Token, error) { return nil, nil },
-			},
+			config: func() *EnterpriseHandlerConfig {
+				c := validEnterpriseHandlerConfig()
+				c.IdPCredentials = nil
+				return c
+			}(),
 			wantError: "IdPCredentials is required",
 		},
 		{
 			name: "invalid IdPCredentials - empty ClientID",
-			config: &EnterpriseHandlerConfig{
-				IdPIssuerURL: "https://idp.example.com",
-				IdPCredentials: &oauthex.ClientCredentials{
-					ClientID: "", // Invalid - empty
-				},
-				MCPAuthServerURL: "https://mcp-auth.example.com",
-				MCPResourceURI:   "https://mcp.example.com",
-				MCPCredentials:   &oauthex.ClientCredentials{ClientID: "id"},
-				IDTokenFetcher:   func(ctx context.Context) (*oauth2.Token, error) { return nil, nil },
-			},
+			config: func() *EnterpriseHandlerConfig {
+				c := validEnterpriseHandlerConfig()
+				c.IdPCredentials = &oauthex.ClientCredentials{ClientID: ""}
+				return c
+			}(),
 			wantError: "invalid IdPCredentials",
 		},
 		{
 			name: "invalid IdPCredentials - empty ClientSecret in ClientSecretAuth",
-			config: &EnterpriseHandlerConfig{
-				IdPIssuerURL: "https://idp.example.com",
-				IdPCredentials: &oauthex.ClientCredentials{
-					ClientID: "idp_client_id",
-					ClientSecretAuth: &oauthex.ClientSecretAuth{
-						ClientSecret: "", // Invalid - empty secret when ClientSecretAuth is set
-					},
-				},
-				MCPAuthServerURL: "https://mcp-auth.example.com",
-				MCPResourceURI:   "https://mcp.example.com",
-				MCPCredentials:   &oauthex.ClientCredentials{ClientID: "id"},
-				IDTokenFetcher:   func(ctx context.Context) (*oauth2.Token, error) { return nil, nil },
-			},
+			config: func() *EnterpriseHandlerConfig {
+				c := validEnterpriseHandlerConfig()
+				c.IdPCredentials = &oauthex.ClientCredentials{
+					ClientID:         "idp_client_id",
+					ClientSecretAuth: &oauthex.ClientSecretAuth{ClientSecret: ""},
+				}
+				return c
+			}(),
 			wantError: "invalid IdPCredentials",
 		},
 		{
 			name: "missing MCPAuthServerURL",
-			config: &EnterpriseHandlerConfig{
-				IdPIssuerURL:     "https://idp.example.com",
-				IdPCredentials:   &oauthex.ClientCredentials{ClientID: "id"},
-				MCPAuthServerURL: "",
-				MCPResourceURI:   "https://mcp.example.com",
-				MCPCredentials:   &oauthex.ClientCredentials{ClientID: "id"},
-				IDTokenFetcher:   func(ctx context.Context) (*oauth2.Token, error) { return nil, nil },
-			},
+			config: func() *EnterpriseHandlerConfig {
+				c := validEnterpriseHandlerConfig()
+				c.MCPAuthServerURL = ""
+				return c
+			}(),
 			wantError: "MCPAuthServerURL is required",
 		},
 		{
 			name: "missing MCPResourceURI",
-			config: &EnterpriseHandlerConfig{
-				IdPIssuerURL:     "https://idp.example.com",
-				IdPCredentials:   &oauthex.ClientCredentials{ClientID: "id"},
-				MCPAuthServerURL: "https://mcp-auth.example.com",
-				MCPResourceURI:   "",
-				MCPCredentials:   &oauthex.ClientCredentials{ClientID: "id"},
-				IDTokenFetcher:   func(ctx context.Context) (*oauth2.Token, error) { return nil, nil },
-			},
+			config: func() *EnterpriseHandlerConfig {
+				c := validEnterpriseHandlerConfig()
+				c.MCPResourceURI = ""
+				return c
+			}(),
 			wantError: "MCPResourceURI is required",
 		},
 		{
 			name: "nil MCPCredentials",
-			config: &EnterpriseHandlerConfig{
-				IdPIssuerURL:     "https://idp.example.com",
-				IdPCredentials:   &oauthex.ClientCredentials{ClientID: "id"},
-				MCPAuthServerURL: "https://mcp-auth.example.com",
-				MCPResourceURI:   "https://mcp.example.com",
-				MCPCredentials:   nil,
-				IDTokenFetcher:   func(ctx context.Context) (*oauth2.Token, error) { return nil, nil },
-			},
+			config: func() *EnterpriseHandlerConfig {
+				c := validEnterpriseHandlerConfig()
+				c.MCPCredentials = nil
+				return c
+			}(),
 			wantError: "MCPCredentials is required",
 		},
 		{
 			name: "invalid MCPCredentials - empty ClientID",
-			config: &EnterpriseHandlerConfig{
-				IdPIssuerURL:     "https://idp.example.com",
-				IdPCredentials:   &oauthex.ClientCredentials{ClientID: "id"},
-				MCPAuthServerURL: "https://mcp-auth.example.com",
-				MCPResourceURI:   "https://mcp.example.com",
-				MCPCredentials: &oauthex.ClientCredentials{
-					ClientID: "", // Invalid - empty
-				},
-				IDTokenFetcher: func(ctx context.Context) (*oauth2.Token, error) { return nil, nil },
-			},
+			config: func() *EnterpriseHandlerConfig {
+				c := validEnterpriseHandlerConfig()
+				c.MCPCredentials = &oauthex.ClientCredentials{ClientID: ""}
+				return c
+			}(),
 			wantError: "invalid MCPCredentials",
 		},
 		{
 			name: "missing IDTokenFetcher",
-			config: &EnterpriseHandlerConfig{
-				IdPIssuerURL:     "https://idp.example.com",
-				IdPCredentials:   &oauthex.ClientCredentials{ClientID: "id"},
-				MCPAuthServerURL: "https://mcp-auth.example.com",
-				MCPResourceURI:   "https://mcp.example.com",
-				MCPCredentials:   &oauthex.ClientCredentials{ClientID: "id"},
-				IDTokenFetcher:   nil,
-			},
+			config: func() *EnterpriseHandlerConfig {
+				c := validEnterpriseHandlerConfig()
+				c.IDTokenFetcher = nil
+				return c
+			}(),
 			wantError: "IDTokenFetcher is required",
 		},
 		{
-			name:      "valid config - public clients (no ClientSecretAuth)",
-			config:    validConfig,
+			name:      "valid public clients",
+			config:    validEnterpriseHandlerConfig(),
 			wantError: "",
 		},
 		{
-			name: "valid config - confidential clients (with ClientSecretAuth)",
-			config: &EnterpriseHandlerConfig{
-				IdPIssuerURL: "https://idp.example.com",
-				IdPCredentials: &oauthex.ClientCredentials{
-					ClientID: "idp_client_id",
-					ClientSecretAuth: &oauthex.ClientSecretAuth{
-						ClientSecret: "idp_secret",
-					},
-				},
-				MCPAuthServerURL: "https://mcp-auth.example.com",
-				MCPResourceURI:   "https://mcp.example.com",
-				MCPCredentials: &oauthex.ClientCredentials{
-					ClientID: "mcp_client_id",
-					ClientSecretAuth: &oauthex.ClientSecretAuth{
-						ClientSecret: "mcp_secret",
-					},
-				},
-				IDTokenFetcher: func(ctx context.Context) (*oauth2.Token, error) {
-					token := &oauth2.Token{
-						AccessToken: "mock_access_token",
-						TokenType:   "Bearer",
-					}
-					return token.WithExtra(map[string]any{"id_token": "mock_id_token"}), nil
-				},
-			},
+			name: "valid confidential clients",
+			config: func() *EnterpriseHandlerConfig {
+				c := validEnterpriseHandlerConfig()
+				c.IdPCredentials.ClientSecretAuth = &oauthex.ClientSecretAuth{ClientSecret: "idp_secret"}
+				c.MCPCredentials.ClientSecretAuth = &oauthex.ClientSecretAuth{ClientSecret: "mcp_secret"}
+				return c
+			}(),
 			wantError: "",
 		},
 	}
@@ -221,45 +165,47 @@ func TestNewEnterpriseHandler_Validation(t *testing.T) {
 	}
 }
 
-// TestEnterpriseHandler_Authorize_E2E tests the complete enterprise authorization flow.
 func TestEnterpriseHandler_Authorize_E2E(t *testing.T) {
-	// Set up IdP (Identity Provider) fake server with token exchange support
-	idpServer := setupIdPServer(t)
+	const idJAGToken = "id-jag-token-from-idp"
 
-	// Set up MCP authorization server with JWT bearer grant support
-	mcpAuthServer := setupMCPAuthServer(t)
-
-	// Create enterprise handler
-	handler, err := NewEnterpriseHandler(&EnterpriseHandlerConfig{
-		IdPIssuerURL: idpServer.URL,
-		IdPCredentials: &oauthex.ClientCredentials{
-			ClientID: "idp_client_id",
-			ClientSecretAuth: &oauthex.ClientSecretAuth{
-				ClientSecret: "idp_secret",
-			},
+	idpServer := oauthtest.NewFakeIdPServer(oauthtest.IdPConfig{
+		PreregisteredClients: map[string]oauthtest.ClientInfo{
+			"idp_client_id": {Secret: "idp_secret"},
 		},
-		MCPAuthServerURL: mcpAuthServer.URL,
-		MCPResourceURI:   "https://mcp.example.com",
-		MCPCredentials: &oauthex.ClientCredentials{
-			ClientID: "mcp_client_id",
-			ClientSecretAuth: &oauthex.ClientSecretAuth{
-				ClientSecret: "mcp_secret",
-			},
-		},
-		MCPScopes: []string{"read", "write"},
-		IDTokenFetcher: func(ctx context.Context) (*oauth2.Token, error) {
-			token := &oauth2.Token{
-				AccessToken: "mock_access_token",
-				TokenType:   "Bearer",
-			}
-			return token.WithExtra(map[string]any{"id_token": "mock_id_token_from_user_login"}), nil
+		TokenExchangeConfig: &oauthtest.TokenExchangeConfig{
+			IDJAGToken: idJAGToken,
 		},
 	})
+	idpServer.Start(t)
+
+	mcpAuthServer := oauthtest.NewFakeAuthorizationServer(oauthtest.Config{
+		RegistrationConfig: &oauthtest.RegistrationConfig{
+			PreregisteredClients: map[string]oauthtest.ClientInfo{
+				"mcp_client_id": {Secret: "mcp_secret"},
+			},
+		},
+		JWTBearerConfig: &oauthtest.JWTBearerConfig{
+			ValidAssertions: []string{idJAGToken},
+		},
+	})
+	mcpAuthServer.Start(t)
+
+	config := validEnterpriseHandlerConfig()
+	config.IdPIssuerURL = idpServer.URL()
+	config.IdPCredentials.ClientSecretAuth = &oauthex.ClientSecretAuth{ClientSecret: "idp_secret"}
+	config.MCPAuthServerURL = mcpAuthServer.URL()
+	config.MCPCredentials.ClientSecretAuth = &oauthex.ClientSecretAuth{ClientSecret: "mcp_secret"}
+	config.MCPScopes = []string{"read", "write"}
+	config.IDTokenFetcher = func(ctx context.Context) (*oauth2.Token, error) {
+		token := &oauth2.Token{AccessToken: "mock_access_token", TokenType: "Bearer"}
+		return token.WithExtra(map[string]any{"id_token": "mock_id_token_from_user_login"}), nil
+	}
+
+	handler, err := NewEnterpriseHandler(config)
 	if err != nil {
 		t.Fatalf("NewEnterpriseHandler failed: %v", err)
 	}
 
-	// Simulate a 401 response from MCP server
 	req := httptest.NewRequest(http.MethodGet, "https://mcp.example.com/api", nil)
 	resp := &http.Response{
 		StatusCode: http.StatusUnauthorized,
@@ -268,12 +214,10 @@ func TestEnterpriseHandler_Authorize_E2E(t *testing.T) {
 		Request:    req,
 	}
 
-	// Perform authorization
 	if err := handler.Authorize(context.Background(), req, resp); err != nil {
 		t.Fatalf("Authorize failed: %v", err)
 	}
 
-	// Verify token source is set
 	tokenSource, err := handler.TokenSource(context.Background())
 	if err != nil {
 		t.Fatalf("TokenSource failed: %v", err)
@@ -282,164 +226,22 @@ func TestEnterpriseHandler_Authorize_E2E(t *testing.T) {
 		t.Fatal("expected token source to be set after authorization")
 	}
 
-	// Verify we can get a token
 	token, err := tokenSource.Token()
 	if err != nil {
 		t.Fatalf("Token() failed: %v", err)
 	}
-	if token.AccessToken != "mcp_access_token_from_jwt_bearer" {
-		t.Errorf("unexpected access token: got %q, want %q",
-			token.AccessToken, "mcp_access_token_from_jwt_bearer")
+	if token.AccessToken != "test_access_token" {
+		t.Errorf("AccessToken = %q, want %q", token.AccessToken, "test_access_token")
 	}
 }
 
-// setupIdPServer creates a fake IdP server that supports token exchange.
-func setupIdPServer(t *testing.T) *httptest.Server {
-	t.Helper()
-	mux := http.NewServeMux()
-
-	var server *httptest.Server
-
-	// OAuth/OIDC metadata endpoint - uses closure to get server URL
-	mux.HandleFunc("/.well-known/oauth-authorization-server", func(w http.ResponseWriter, r *http.Request) {
-		w.Header().Set("Content-Type", "application/json")
-		json.NewEncoder(w).Encode(map[string]any{
-			"issuer":                           server.URL,
-			"token_endpoint":                   server.URL + "/token",
-			"authorization_endpoint":           server.URL + "/authorize",
-			"code_challenge_methods_supported": []string{"S256"},
-		})
-	})
-
-	// Token endpoint - supports token exchange (RFC 8693)
-	mux.HandleFunc("/token", func(w http.ResponseWriter, r *http.Request) {
-		if err := r.ParseForm(); err != nil {
-			http.Error(w, "failed to parse form", http.StatusBadRequest)
-			return
-		}
-
-		grantType := r.Form.Get("grant_type")
-		if grantType != oauthex.GrantTypeTokenExchange {
-			http.Error(w, fmt.Sprintf("unsupported grant_type: %s", grantType), http.StatusBadRequest)
-			return
-		}
-
-		// Verify client authentication
-		clientID := r.Form.Get("client_id")
-		clientSecret := r.Form.Get("client_secret")
-		if clientID != "idp_client_id" || clientSecret != "idp_secret" {
-			http.Error(w, "invalid client credentials", http.StatusUnauthorized)
-			return
-		}
-
-		// Verify token exchange parameters
-		if r.Form.Get("requested_token_type") != oauthex.TokenTypeIDJAG {
-			http.Error(w, "invalid requested_token_type", http.StatusBadRequest)
-			return
-		}
-		if r.Form.Get("subject_token_type") != oauthex.TokenTypeIDToken {
-			http.Error(w, "invalid subject_token_type", http.StatusBadRequest)
-			return
-		}
-		if r.Form.Get("subject_token") == "" {
-			http.Error(w, "missing subject_token", http.StatusBadRequest)
-			return
-		}
-
-		// Return ID-JAG (Identity Assertion JWT Authorization Grant)
-		w.Header().Set("Content-Type", "application/json")
-		json.NewEncoder(w).Encode(map[string]any{
-			"access_token":      "id-jag-token-from-idp",
-			"issued_token_type": oauthex.TokenTypeIDJAG,
-			"token_type":        "N_A",
-			"expires_in":        300,
-			"scope":             "read write",
-		})
-	})
-
-	server = httptest.NewServer(mux)
-	t.Cleanup(server.Close)
-
-	return server
-}
-
-// setupMCPAuthServer creates a fake MCP authorization server that supports JWT bearer grant.
-func setupMCPAuthServer(t *testing.T) *httptest.Server {
-	t.Helper()
-	mux := http.NewServeMux()
-
-	var server *httptest.Server
-
-	// OAuth metadata endpoint - uses closure to get server URL
-	mux.HandleFunc("/.well-known/oauth-authorization-server", func(w http.ResponseWriter, r *http.Request) {
-		w.Header().Set("Content-Type", "application/json")
-		json.NewEncoder(w).Encode(map[string]any{
-			"issuer":                           server.URL,
-			"token_endpoint":                   server.URL + "/token",
-			"code_challenge_methods_supported": []string{"S256"},
-		})
-	})
-
-	// Token endpoint - supports JWT bearer grant (RFC 7523)
-	mux.HandleFunc("/token", func(w http.ResponseWriter, r *http.Request) {
-		if err := r.ParseForm(); err != nil {
-			http.Error(w, "failed to parse form", http.StatusBadRequest)
-			return
-		}
-
-		grantType := r.Form.Get("grant_type")
-		if grantType != "urn:ietf:params:oauth:grant-type:jwt-bearer" {
-			http.Error(w, fmt.Sprintf("unsupported grant_type: %s", grantType), http.StatusBadRequest)
-			return
-		}
-
-		// Verify client authentication
-		clientID := r.Form.Get("client_id")
-		clientSecret := r.Form.Get("client_secret")
-		if clientID != "mcp_client_id" || clientSecret != "mcp_secret" {
-			http.Error(w, "invalid client credentials", http.StatusUnauthorized)
-			return
-		}
-
-		// Verify assertion (ID-JAG)
-		assertion := r.Form.Get("assertion")
-		if assertion != "id-jag-token-from-idp" {
-			http.Error(w, "invalid assertion", http.StatusBadRequest)
-			return
-		}
-
-		// Return access token
-		w.Header().Set("Content-Type", "application/json")
-		json.NewEncoder(w).Encode(map[string]any{
-			"access_token": "mcp_access_token_from_jwt_bearer",
-			"token_type":   "Bearer",
-			"expires_in":   3600,
-			"scope":        "read write",
-		})
-	})
-
-	server = httptest.NewServer(mux)
-	t.Cleanup(server.Close)
-
-	return server
-}
-
-// TestEnterpriseHandler_Authorize_IDTokenFetcherError tests error handling when IDTokenFetcher fails.
 func TestEnterpriseHandler_Authorize_IDTokenFetcherError(t *testing.T) {
-	handler, err := NewEnterpriseHandler(&EnterpriseHandlerConfig{
-		IdPIssuerURL: "https://idp.example.com",
-		IdPCredentials: &oauthex.ClientCredentials{
-			ClientID: "idp_client_id",
-		},
-		MCPAuthServerURL: "https://mcp-auth.example.com",
-		MCPResourceURI:   "https://mcp.example.com",
-		MCPCredentials: &oauthex.ClientCredentials{
-			ClientID: "mcp_client_id",
-		},
-		IDTokenFetcher: func(ctx context.Context) (*oauth2.Token, error) {
-			return nil, fmt.Errorf("user cancelled login")
-		},
-	})
+	config := validEnterpriseHandlerConfig()
+	config.IDTokenFetcher = func(ctx context.Context) (*oauth2.Token, error) {
+		return nil, fmt.Errorf("user cancelled login")
+	}
+
+	handler, err := NewEnterpriseHandler(config)
 	if err != nil {
 		t.Fatalf("NewEnterpriseHandler failed: %v", err)
 	}
@@ -457,30 +259,12 @@ func TestEnterpriseHandler_Authorize_IDTokenFetcherError(t *testing.T) {
 		t.Fatal("expected error from Authorize, got nil")
 	}
 	if !strings.Contains(err.Error(), "failed to obtain ID token") {
-		t.Errorf("expected error about ID token, got: %v", err)
+		t.Errorf("error = %v, want error containing %q", err, "failed to obtain ID token")
 	}
 }
 
-// TestEnterpriseHandler_TokenSource_BeforeAuthorization tests TokenSource before authorization.
 func TestEnterpriseHandler_TokenSource_BeforeAuthorization(t *testing.T) {
-	handler, err := NewEnterpriseHandler(&EnterpriseHandlerConfig{
-		IdPIssuerURL: "https://idp.example.com",
-		IdPCredentials: &oauthex.ClientCredentials{
-			ClientID: "idp_client_id",
-		},
-		MCPAuthServerURL: "https://mcp-auth.example.com",
-		MCPResourceURI:   "https://mcp.example.com",
-		MCPCredentials: &oauthex.ClientCredentials{
-			ClientID: "mcp_client_id",
-		},
-		IDTokenFetcher: func(ctx context.Context) (*oauth2.Token, error) {
-			token := &oauth2.Token{
-				AccessToken: "mock_access_token",
-				TokenType:   "Bearer",
-			}
-			return token.WithExtra(map[string]any{"id_token": "mock_id_token"}), nil
-		},
-	})
+	handler, err := NewEnterpriseHandler(validEnterpriseHandlerConfig())
 	if err != nil {
 		t.Fatalf("NewEnterpriseHandler failed: %v", err)
 	}