cneicy
ceddfd3d6b
Sync production source including routing, billing, Astra detection and public uptime. Update custom build instructions and license packaging.
103 lines
3.9 KiB
Go
103 lines
3.9 KiB
Go
package middleware
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"net/http"
|
|
"sync/atomic"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/QuantumNous/new-api/common"
|
|
"github.com/QuantumNous/new-api/constant"
|
|
relaycommon "github.com/QuantumNous/new-api/relay/common"
|
|
"github.com/gin-gonic/gin"
|
|
"github.com/stretchr/testify/assert"
|
|
"github.com/stretchr/testify/require"
|
|
)
|
|
|
|
func TestModelRedisRateLimitUsesUTCRegardlessOfLocalTimezone(t *testing.T) {
|
|
redisServer, redisClient := useRateLimitMiniRedis(t)
|
|
previousLocation := time.Local
|
|
time.Local = time.FixedZone("test-utc-plus-eight", 8*60*60)
|
|
t.Cleanup(func() { time.Local = previousLocation })
|
|
|
|
ctx := context.Background()
|
|
recordKey := "rateLimit:model-utc-record"
|
|
recordRedisRequest(ctx, redisClient, recordKey, 2)
|
|
recorded, err := redisClient.LIndex(ctx, recordKey, 0).Result()
|
|
require.NoError(t, err)
|
|
recordedAt, err := time.Parse(modelRateLimitTimeFormat, recorded)
|
|
require.NoError(t, err)
|
|
assert.WithinDuration(t, time.Now().UTC(), recordedAt, 2*time.Second)
|
|
|
|
checkKey := "rateLimit:model-utc-check"
|
|
withinWindow := time.Now().UTC().Add(-30 * time.Second).Format(modelRateLimitTimeFormat)
|
|
_, err = redisServer.Push(checkKey, withinWindow, withinWindow)
|
|
require.NoError(t, err)
|
|
allowed, err := checkRedisRateLimit(ctx, redisClient, checkKey, 2, 60)
|
|
require.NoError(t, err)
|
|
assert.False(t, allowed, "an existing UTC timestamp inside the window must remain limited on a non-UTC host")
|
|
}
|
|
|
|
var modelRateLimitTestUsers atomic.Int64
|
|
|
|
func TestModelRateLimitStreamFailuresDoNotConsumeSuccessLimit(t *testing.T) {
|
|
for _, backend := range []string{"memory", "redis"} {
|
|
for _, totalLimit := range []int{0, 2} {
|
|
t.Run(fmt.Sprintf("%s/total=%d", backend, totalLimit), func(t *testing.T) {
|
|
userID := 7200000 + int(modelRateLimitTestUsers.Add(1))
|
|
handler := memoryRateLimitHandler(60, totalLimit, 1)
|
|
if backend == "redis" {
|
|
useRateLimitMiniRedis(t)
|
|
handler = redisRateLimitHandler(60, totalLimit, 1)
|
|
}
|
|
router := gin.New()
|
|
router.GET("/:outcome", func(c *gin.Context) { c.Set("id", userID) }, handler, func(c *gin.Context) {
|
|
status := relaycommon.NewStreamStatus()
|
|
if c.Param("outcome") == "failed" {
|
|
status.MarkFailed("server_error", "", 0)
|
|
} else {
|
|
status.MarkCompleted()
|
|
}
|
|
common.SetContextKey(c, constant.ContextKeyResponseStreamStatus, status)
|
|
c.Status(http.StatusOK)
|
|
})
|
|
assert.Equal(t, http.StatusOK, performRateLimitRequest(router, "/failed", "127.0.0.1:1000").Code)
|
|
if totalLimit > 0 {
|
|
assert.Equal(t, http.StatusOK, performRateLimitRequest(router, "/failed", "127.0.0.1:1000").Code)
|
|
} else {
|
|
assert.Equal(t, http.StatusOK, performRateLimitRequest(router, "/completed", "127.0.0.1:1000").Code)
|
|
}
|
|
assert.Equal(t, http.StatusTooManyRequests, performRateLimitRequest(router, "/completed", "127.0.0.1:1000").Code)
|
|
})
|
|
}
|
|
}
|
|
}
|
|
|
|
func TestModelMemoryRateLimitReservesConcurrentSuccessAdmission(t *testing.T) {
|
|
userID := 7200000 + int(modelRateLimitTestUsers.Add(1))
|
|
entered, release, finished := make(chan struct{}), make(chan struct{}), make(chan struct{})
|
|
router := gin.New()
|
|
router.GET("/:outcome", func(c *gin.Context) { c.Set("id", userID) }, memoryRateLimitHandler(60, 0, 1), func(c *gin.Context) {
|
|
if c.Param("outcome") == "failed" {
|
|
close(entered)
|
|
<-release
|
|
status := relaycommon.NewStreamStatus()
|
|
status.MarkFailed("server_error", "", 0)
|
|
common.SetContextKey(c, constant.ContextKeyResponseStreamStatus, status)
|
|
}
|
|
c.Status(http.StatusOK)
|
|
})
|
|
go func() {
|
|
defer close(finished)
|
|
assert.Equal(t, http.StatusOK, performRateLimitRequest(router, "/failed", "127.0.0.1:1000").Code)
|
|
}()
|
|
<-entered
|
|
assert.Equal(t, http.StatusTooManyRequests, performRateLimitRequest(router, "/completed", "127.0.0.1:1000").Code)
|
|
close(release)
|
|
<-finished
|
|
assert.Equal(t, http.StatusOK, performRateLimitRequest(router, "/completed", "127.0.0.1:1000").Code)
|
|
assert.Equal(t, http.StatusTooManyRequests, performRateLimitRequest(router, "/completed", "127.0.0.1:1000").Code)
|
|
}
|