优化结构
This commit is contained in:
@@ -1,70 +0,0 @@
|
||||
package middleware
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"kra/internal/biz"
|
||||
"kra/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())
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user