71 lines
3.1 KiB
Go
71 lines
3.1 KiB
Go
package handler
|
|
|
|
import (
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"testing"
|
|
"time"
|
|
|
|
"git.ilapage.cn/OPC/chorus/internal/platform/ratelimit"
|
|
"git.ilapage.cn/OPC/chorus/portal/service"
|
|
"git.ilapage.cn/OPC/chorus/portal/session"
|
|
"github.com/gin-gonic/gin"
|
|
)
|
|
|
|
func TestPortalSubmissionRateLimit(t *testing.T) {
|
|
gin.SetMode(gin.TestMode)
|
|
now := time.Unix(1_700_000_000, 0).UTC()
|
|
handler := &Handler{
|
|
rateLimiter: ratelimit.New(), userRatePolicy: ratelimit.Policy{Capacity: 1, Window: time.Minute},
|
|
now: func() time.Time { return now },
|
|
}
|
|
router := gin.New()
|
|
router.POST("/api/generations/text", func(c *gin.Context) {
|
|
c.Set("session", session.State{UserID: 7})
|
|
c.Next()
|
|
}, handler.limitPortalSubmission, func(c *gin.Context) { c.Status(http.StatusNoContent) })
|
|
|
|
if response := rateRequest(router, http.MethodPost, "/api/generations/text"); response.Code != http.StatusNoContent {
|
|
t.Fatalf("first request = %d %s", response.Code, response.Body.String())
|
|
}
|
|
response := rateRequest(router, http.MethodPost, "/api/generations/text")
|
|
if response.Code != http.StatusTooManyRequests || response.Header().Get("Retry-After") != "60" || response.Body.String() != `{"error":{"code":"rate_limited","message":"request rate limit exceeded"}}` {
|
|
t.Fatalf("limited request = %d retry=%q body=%s", response.Code, response.Header().Get("Retry-After"), response.Body.String())
|
|
}
|
|
}
|
|
|
|
func TestOpenAPISubmissionChecksUserAndKeyAtomically(t *testing.T) {
|
|
gin.SetMode(gin.TestMode)
|
|
now := time.Unix(1_700_000_000, 0).UTC()
|
|
handler := &Handler{
|
|
rateLimiter: ratelimit.New(), userRatePolicy: ratelimit.Policy{Capacity: 1, Window: time.Minute},
|
|
apiKeyRatePolicy: ratelimit.Policy{Capacity: 2, Window: time.Minute}, now: func() time.Time { return now },
|
|
}
|
|
router := gin.New()
|
|
router.Use(func(c *gin.Context) {
|
|
c.Set(openAPIPrincipalContextKey, service.APIPrincipal{UserID: 7, APIKeyID: 11})
|
|
c.Next()
|
|
})
|
|
router.POST("/openapi/v1/generations/text", handler.limitOpenAPIRequest, func(c *gin.Context) { c.Status(http.StatusNoContent) })
|
|
router.GET("/openapi/v1/generations", handler.limitOpenAPIRequest, func(c *gin.Context) { c.Status(http.StatusNoContent) })
|
|
|
|
if response := rateRequest(router, http.MethodPost, "/openapi/v1/generations/text"); response.Code != http.StatusNoContent {
|
|
t.Fatalf("first submission = %d", response.Code)
|
|
}
|
|
if response := rateRequest(router, http.MethodPost, "/openapi/v1/generations/text"); response.Code != http.StatusTooManyRequests {
|
|
t.Fatalf("user-limited submission = %d", response.Code)
|
|
}
|
|
if response := rateRequest(router, http.MethodGet, "/openapi/v1/generations"); response.Code != http.StatusNoContent {
|
|
t.Fatalf("rejected combined check consumed key bucket: %d", response.Code)
|
|
}
|
|
if response := rateRequest(router, http.MethodGet, "/openapi/v1/generations"); response.Code != http.StatusTooManyRequests {
|
|
t.Fatalf("API key total request limit = %d", response.Code)
|
|
}
|
|
}
|
|
|
|
func rateRequest(router http.Handler, method, path string) *httptest.ResponseRecorder {
|
|
response := httptest.NewRecorder()
|
|
router.ServeHTTP(response, httptest.NewRequest(method, path, nil))
|
|
return response
|
|
}
|