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") } }