package middleware import ( "context" "errors" "net/http" "net/http/httptest" "strings" "testing" "time" "kra/app/system/internal/biz" "kra/app/system/internal/service" "github.com/gin-gonic/gin" ) type rateLimitSecurityRepo struct{} func (rateLimitSecurityRepo) SecurityConfig(context.Context) (*biz.SecurityConfig, error) { return &biz.SecurityConfig{LimitEnable: true, LimitWindow: 60, LimitCount: 1}, nil } func (rateLimitSecurityRepo) SaveSecurityConfig(context.Context, *biz.SecurityConfig) error { return nil } func (rateLimitSecurityRepo) BackfillPasswordUpdatedAt(context.Context, time.Time) error { return nil } type rateLimitCache struct{ err error } func (rateLimitCache) Get(context.Context, string) (string, bool, error) { return "", false, nil } func (rateLimitCache) Set(context.Context, string, string, time.Duration) error { return nil } func (rateLimitCache) Delete(context.Context, string) error { return nil } func (c rateLimitCache) Increment(context.Context, string, time.Duration) (int64, error) { return 2, c.err } func TestSecurityRateLimitMatchesResponseContract(t *testing.T) { gin.SetMode(gin.TestMode) settings := service.NewSecurityService(biz.NewSecurityUsecase(rateLimitSecurityRepo{}, rateLimitCache{}, nil, nil)) engine := gin.New() engine.Use(SecurityRateLimit(settings)) engine.POST("/base/login", func(c *gin.Context) { c.Status(http.StatusNoContent) }) response := httptest.NewRecorder() engine.ServeHTTP(response, httptest.NewRequest(http.MethodPost, "/base/login", nil)) if response.Code != http.StatusOK { t.Fatalf("status = %d, body=%s", response.Code, response.Body.String()) } body := response.Body.String() if !strings.Contains(body, `"code":7`) || !strings.Contains(body, `"msg":"请求太过频繁,请稍后再试"`) || strings.Contains(body, `"data"`) { t.Fatalf("unexpected rate-limit response: %s", body) } } func TestSecurityRateLimitFailsClosedWhenCacheIsUnavailable(t *testing.T) { gin.SetMode(gin.TestMode) settings := service.NewSecurityService(biz.NewSecurityUsecase(rateLimitSecurityRepo{}, rateLimitCache{err: errors.New("cache unavailable")}, nil, nil)) engine := gin.New() called := false engine.Use(SecurityRateLimit(settings)) engine.POST("/base/login", func(c *gin.Context) { called = true }) response := httptest.NewRecorder() engine.ServeHTTP(response, httptest.NewRequest(http.MethodPost, "/base/login", nil)) if response.Code != http.StatusServiceUnavailable || called { t.Fatalf("status=%d called=%v body=%s", response.Code, called, response.Body.String()) } }