MyAiDesk/GovAI/server/internal/middleware/middleware_test.go

package middleware

import (
	"context"
	"net/http"
	"net/http/httptest"
	"testing"

	"github.com/google/uuid"
)

func TestGetUserID(t *testing.T) {
	userID := uuid.New()
	ctx := context.WithValue(context.Background(), UserIDKey, userID)
	if got := GetUserID(ctx); got != userID {
		t.Errorf("GetUserID() = %v, want %v", got, userID)
	}

	ctx = context.Background()
	if got := GetUserID(ctx); got != uuid.Nil {
		t.Errorf("GetUserID() from empty ctx = %v, want uuid.Nil", got)
	}
}

func TestGetRole(t *testing.T) {
	ctx := context.WithValue(context.Background(), RoleKey, "admin")
	if got := GetRole(ctx); got != "admin" {
		t.Errorf("GetRole() = %v, want admin", got)
	}

	ctx = context.Background()
	if got := GetRole(ctx); got != "" {
		t.Errorf("GetRole() from empty ctx = %v, want empty string", got)
	}
}

func TestRequireRole_UserRole(t *testing.T) {
	mux := http.NewServeMux()
	handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusOK)
	})

	wrapped := RequireRole("admin")(handler)
	mux.Handle("/admin", wrapped)

	// user 角色 → 403
	userCtx := context.WithValue(context.Background(), RoleKey, "user")
	req := httptest.NewRequest(http.MethodGet, "/admin", nil).WithContext(userCtx)
	rr := httptest.NewRecorder()
	mux.ServeHTTP(rr, req)
	if rr.Code != http.StatusForbidden {
		t.Errorf("user role: got %d, want %d", rr.Code, http.StatusForbidden)
	}

	// admin 角色 → 200
	adminCtx := context.WithValue(context.Background(), RoleKey, "admin")
	req = httptest.NewRequest(http.MethodGet, "/admin", nil).WithContext(adminCtx)
	rr = httptest.NewRecorder()
	mux.ServeHTTP(rr, req)
	if rr.Code != http.StatusOK {
		t.Errorf("admin role: got %d, want %d", rr.Code, http.StatusOK)
	}

	// 无角色 → 403
	req = httptest.NewRequest(http.MethodGet, "/admin", nil)
	rr = httptest.NewRecorder()
	mux.ServeHTTP(rr, req)
	if rr.Code != http.StatusForbidden {
		t.Errorf("no role: got %d, want %d", rr.Code, http.StatusForbidden)
	}
}

func TestRequireSuperAdmin(t *testing.T) {
	mux := http.NewServeMux()
	handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusOK)
	})

	wrapped := RequireSuperAdmin(handler)
	mux.Handle("/platform", wrapped)

	// super_admin → 200
	superAdminCtx := context.WithValue(context.Background(), RoleKey, "super_admin")
	req := httptest.NewRequest(http.MethodGet, "/platform", nil).WithContext(superAdminCtx)
	rr := httptest.NewRecorder()
	mux.ServeHTTP(rr, req)
	if rr.Code != http.StatusOK {
		t.Errorf("super_admin: got %d, want %d", rr.Code, http.StatusOK)
	}

	// admin → 403
	adminCtx := context.WithValue(context.Background(), RoleKey, "admin")
	req = httptest.NewRequest(http.MethodGet, "/platform", nil).WithContext(adminCtx)
	rr = httptest.NewRecorder()
	mux.ServeHTTP(rr, req)
	if rr.Code != http.StatusForbidden {
		t.Errorf("admin: got %d, want %d", rr.Code, http.StatusForbidden)
	}
}

func TestAuditLog_NilPool(t *testing.T) {
	req := httptest.NewRequest(http.MethodPost, "/test", nil)
	rr := httptest.NewRecorder()

	fn := AuditLog(nil)
	var called bool
	fn(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		called = true
	})).ServeHTTP(rr, req)

	if !called {
		t.Error("AuditLog middleware did not call next handler")
	}
}

func TestRateLimit_NilRedis(t *testing.T) {
	mux := http.NewServeMux()
	handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusOK)
	})

	wrapped := RateLimit(nil, 5, 0)(handler)
	mux.Handle("/test", wrapped)

	req := httptest.NewRequest(http.MethodGet, "/test", nil)
	rr := httptest.NewRecorder()
	mux.ServeHTTP(rr, req)

	if rr.Code == http.StatusTooManyRequests {
		t.Error("RateLimit should bypass when Redis is unavailable")
	}
}