Edu/services/api-gateway/internal/proxy/proxy_rewrite_test.go

package proxy

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

	"github.com/gin-gonic/gin"
)

// startGateway 启动一个真实 httptest.Server 作为 api-gateway，返回其 URL。
// 使用真实 server 是因为 httputil.ReverseProxy.ServeHTTP 会调用
// ResponseWriter.CloseNotify()，而 httptest.ResponseRecorder 未实现该接口。
func startGateway(t *testing.T, handler gin.HandlerFunc) *httptest.Server {
	t.Helper()
	gin.SetMode(gin.TestMode)
	r := gin.New()
	// 只注册 graphql 具体路由，避免与通配符 *path 冲突
	r.POST("/api/admin/graphql", handler)
	r.GET("/api/admin/graphql", handler)
	return httptest.NewServer(r)
}

// TestNewProxyRewrite_StripsPrefix 验证 NewProxyRewrite 调用自定义 rewriter 后下游收到的路径正确。
// 模拟 admin-portal 场景：/api/admin/graphql → /graphql
func TestNewProxyRewrite_StripsPrefix(t *testing.T) {
	// 下游服务器：记录收到的路径
	var receivedPath string
	downstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		receivedPath = r.URL.Path
		w.WriteHeader(http.StatusOK)
		_, _ = w.Write([]byte("ok"))
	}))
	defer downstream.Close()

	rewriter := func(p string) string {
		return strings.TrimPrefix(p, "/api/admin")
	}
	p, err := NewProxyRewrite(downstream.URL, rewriter)
	if err != nil {
		t.Fatalf("NewProxyRewrite 失败: %v", err)
	}

	gateway := startGateway(t, ProxyHandler(p))
	defer gateway.Close()

	// 通过真实 HTTP 客户端发起请求
	req, _ := http.NewRequest(http.MethodPost, gateway.URL+"/api/admin/graphql", strings.NewReader(`{"query":"{}"}`))
	req.Header.Set("Content-Type", "application/json")
	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		t.Fatalf("请求 gateway 失败: %v", err)
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		t.Fatalf("期望 200，实际 %d", resp.StatusCode)
	}
	if receivedPath != "/graphql" {
		t.Fatalf("下游收到路径应为 /graphql，实际 %s", receivedPath)
	}
}

// TestNewProxyRewrite_PreservesQuery 验证 NewProxyRewrite 保留查询参数。
func TestNewProxyRewrite_PreservesQuery(t *testing.T) {
	var receivedQuery string
	downstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		receivedQuery = r.URL.RawQuery
		w.WriteHeader(http.StatusOK)
	}))
	defer downstream.Close()

	rewriter := func(p string) string {
		return strings.TrimPrefix(p, "/api/admin")
	}
	p, err := NewProxyRewrite(downstream.URL, rewriter)
	if err != nil {
		t.Fatalf("NewProxyRewrite 失败: %v", err)
	}

	gateway := startGateway(t, ProxyHandler(p))
	defer gateway.Close()

	req, _ := http.NewRequest(http.MethodGet, gateway.URL+"/api/admin/graphql?operation=adminUsers", nil)
	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		t.Fatalf("请求 gateway 失败: %v", err)
	}
	defer resp.Body.Close()

	if receivedQuery != "operation=adminUsers" {
		t.Fatalf("下游应收到查询参数 operation=adminUsers，实际 %s", receivedQuery)
	}
}

// TestNewProxyRewrite_InvalidURL 验证 NewProxyRewrite 对无效 URL 报错。
func TestNewProxyRewrite_InvalidURL(t *testing.T) {
	rewriter := func(p string) string { return p }
	_, err := NewProxyRewrite("://invalid", rewriter)
	if err == nil {
		t.Fatal("期望无效 URL 报错，实际返回 nil")
	}
}

// TestNewProxyRewrite_ForwardsBody 验证 NewProxyRewrite 转发请求体。
func TestNewProxyRewrite_ForwardsBody(t *testing.T) {
	var receivedBody string
	downstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		body, _ := io.ReadAll(r.Body)
		receivedBody = string(body)
		w.WriteHeader(http.StatusOK)
	}))
	defer downstream.Close()

	rewriter := func(p string) string {
		return strings.TrimPrefix(p, "/api/admin")
	}
	p, err := NewProxyRewrite(downstream.URL, rewriter)
	if err != nil {
		t.Fatalf("NewProxyRewrite 失败: %v", err)
	}

	gateway := startGateway(t, ProxyHandler(p))
	defer gateway.Close()

	body := `{"query":"query { adminUsers { id } }"}`
	req, _ := http.NewRequest(http.MethodPost, gateway.URL+"/api/admin/graphql", strings.NewReader(body))
	req.Header.Set("Content-Type", "application/json")
	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		t.Fatalf("请求 gateway 失败: %v", err)
	}
	defer resp.Body.Close()

	if receivedBody != body {
		t.Fatalf("下游应收到完整请求体，实际 %s", receivedBody)
	}
}

// TestNewProxyRewrite_ForwardsHeaders 验证 NewProxyRewrite 转发请求头。
func TestNewProxyRewrite_ForwardsHeaders(t *testing.T) {
	var receivedAuth string
	downstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		receivedAuth = r.Header.Get("x-user-roles")
		w.WriteHeader(http.StatusOK)
	}))
	defer downstream.Close()

	rewriter := func(p string) string {
		return strings.TrimPrefix(p, "/api/admin")
	}
	p, err := NewProxyRewrite(downstream.URL, rewriter)
	if err != nil {
		t.Fatalf("NewProxyRewrite 失败: %v", err)
	}

	gateway := startGateway(t, ProxyHandler(p))
	defer gateway.Close()

	req, _ := http.NewRequest(http.MethodPost, gateway.URL+"/api/admin/graphql", nil)
	req.Header.Set("x-user-roles", "admin")
	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		t.Fatalf("请求 gateway 失败: %v", err)
	}
	defer resp.Body.Close()

	if receivedAuth != "admin" {
		t.Fatalf("下游应收到 x-user-roles=admin，实际 %s", receivedAuth)
	}
}