liuyupeng/sub2api
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Sub2API v1.0 - AI API 网关（二开初始版本，基于上游 Wei-Shaw/sub2api）
Release / update-version (push) Has been cancelled
Details
Release / build-frontend (push) Has been cancelled
Details
Release / release (push) Has been cancelled
Details
Release / sync-version-file (push) Has been cancelled
Details
CI / shell (push) Canceled after 0s
Details
CI / test (push) Canceled after 0s
Details
CI / frontend (push) Canceled after 0s
Details
CI / golangci-lint (push) Canceled after 0s
Details
Security Scan / backend-security (push) Canceled after 0s
Details
Security Scan / frontend-security (push) Canceled after 0s
Details

Browse Source
This commit is contained in:
李建琦
2026-08-21 18:30:13 +08:00
commit 6d655c9903
3584 changed files with 1270640 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/internal/service/model_rate_limit_test.go
+545
View File
@@ -0,0 +1,545 @@
package service
import (
"context"
"testing"
"time"
"github.com/Wei-Shaw/sub2api/internal/pkg/ctxkey"
"github.com/stretchr/testify/require"
)
func TestIsModelRateLimited(t *testing.T) {
now := time.Now()
future := now.Add(10 * time.Minute).Format(time.RFC3339)
past := now.Add(-10 * time.Minute).Format(time.RFC3339)
tests := []struct {
name string
account *Account
requestedModel string
expected bool
}{
{
name: "official model ID hit - claude-sonnet-4-5",
account: &Account{
Extra: map[string]any{
modelRateLimitsKey: map[string]any{
"claude-sonnet-4-5": map[string]any{
"rate_limit_reset_at": future,
},
},
},
},
requestedModel: "claude-sonnet-4-5",
expected: true,
},
{
name: "official model ID hit via mapping - request claude-3-5-sonnet, mapped to claude-sonnet-4-5",
account: &Account{
Credentials: map[string]any{
"model_mapping": map[string]any{
"claude-3-5-sonnet": "claude-sonnet-4-5",
},
},
Extra: map[string]any{
modelRateLimitsKey: map[string]any{
"claude-sonnet-4-5": map[string]any{
"rate_limit_reset_at": future,
},
},
},
},
requestedModel: "claude-3-5-sonnet",
expected: true,
},
{
name: "no rate limit - expired",
account: &Account{
Extra: map[string]any{
modelRateLimitsKey: map[string]any{
"claude-sonnet-4-5": map[string]any{
"rate_limit_reset_at": past,
},
},
},
},
requestedModel: "claude-sonnet-4-5",
expected: false,
},
{
name: "no rate limit - no matching key",
account: &Account{
Extra: map[string]any{
modelRateLimitsKey: map[string]any{
"gemini-3-flash": map[string]any{
"rate_limit_reset_at": future,
},
},
},
},
requestedModel: "claude-sonnet-4-5",
expected: false,
},
{
name: "no rate limit - unsupported model",
account: &Account{},
requestedModel: "gpt-4",
expected: false,
},
{
name: "no rate limit - empty model",
account: &Account{},
requestedModel: "",
expected: false,
},
{
name: "gemini model hit",
account: &Account{
Extra: map[string]any{
modelRateLimitsKey: map[string]any{
"gemini-3-pro-high": map[string]any{
"rate_limit_reset_at": future,
},
},
},
},
requestedModel: "gemini-3-pro-high",
expected: true,
},
{
name: "antigravity platform - gemini-3-pro-preview mapped to gemini-3-pro-high",
account: &Account{
Platform: PlatformAntigravity,
Extra: map[string]any{
modelRateLimitsKey: map[string]any{
"gemini-3-pro-high": map[string]any{
"rate_limit_reset_at": future,
},
},
},
},
requestedModel: "gemini-3-pro-preview",
expected: true,
},
{
name: "antigravity platform - gemini family rate limit blocks mapped preview",
account: &Account{
Platform: PlatformAntigravity,
Extra: map[string]any{
modelRateLimitsKey: map[string]any{
antigravityGeminiModelRateLimitKey: map[string]any{
"rate_limit_reset_at": future,
},
},
},
},
requestedModel: "gemini-3-pro-preview",
expected: true,
},
{
name: "antigravity platform - gemini family rate limit does not block claude",
account: &Account{
Platform: PlatformAntigravity,
Extra: map[string]any{
modelRateLimitsKey: map[string]any{
antigravityGeminiModelRateLimitKey: map[string]any{
"rate_limit_reset_at": future,
},
},
},
},
requestedModel: "claude-sonnet-4-5",
expected: false,
},
{
name: "non-antigravity platform - gemini-3-pro-preview NOT mapped",
account: &Account{
Platform: PlatformGemini,
Extra: map[string]any{
modelRateLimitsKey: map[string]any{
"gemini-3-pro-high": map[string]any{
"rate_limit_reset_at": future,
},
},
},
},
requestedModel: "gemini-3-pro-preview",
expected: false, // gemini 平台不走 antigravity 映射
},
{
name: "antigravity platform - claude-opus-4-5-thinking mapped to opus-4-6-thinking",
account: &Account{
Platform: PlatformAntigravity,
Extra: map[string]any{
modelRateLimitsKey: map[string]any{
"claude-opus-4-6-thinking": map[string]any{
"rate_limit_reset_at": future,
},
},
},
},
requestedModel: "claude-opus-4-5-thinking",
expected: true,
},
{
name: "no scope fallback - claude_sonnet should not match",
account: &Account{
Extra: map[string]any{
modelRateLimitsKey: map[string]any{
"claude_sonnet": map[string]any{
"rate_limit_reset_at": future,
},
},
},
},
requestedModel: "claude-3-5-sonnet-20241022",
expected: false,
},
{
name: "openai image generation family key blocks image model",
account: &Account{
Platform: PlatformOpenAI,
Extra: map[string]any{
modelRateLimitsKey: map[string]any{
openAIImageGenerationRateLimitKey: map[string]any{
"rate_limit_reset_at": future,
},
},
},
},
requestedModel: "gpt-image-2",
expected: true,
},
{
name: "openai image generation family key does not block text model",
account: &Account{
Platform: PlatformOpenAI,
Extra: map[string]any{
modelRateLimitsKey: map[string]any{
openAIImageGenerationRateLimitKey: map[string]any{
"rate_limit_reset_at": future,
},
},
},
},
requestedModel: "gpt-5.4",
expected: false,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
result := tt.account.isModelRateLimitedWithContext(context.Background(), tt.requestedModel)
if result != tt.expected {
t.Errorf("isModelRateLimited(%q) = %v, want %v", tt.requestedModel, result, tt.expected)
}
})
}
}
func TestIsModelRateLimited_OpenAIImageGenerationIntentBlocksTextModelImageTool(t *testing.T) {
future := time.Now().Add(10 * time.Minute).Format(time.RFC3339)
account := &Account{
Platform: PlatformOpenAI,
Extra: map[string]any{
modelRateLimitsKey: map[string]any{
openAIImageGenerationRateLimitKey: map[string]any{
"rate_limit_reset_at": future,
},
},
},
}
require.False(t, account.isModelRateLimitedWithContext(context.Background(), "gpt-5.4"))
require.True(t, account.isModelRateLimitedWithContext(WithOpenAIImageGenerationIntent(context.Background()), "gpt-5.4"))
}
func TestIsModelRateLimited_Antigravity_ThinkingAffectsModelKey(t *testing.T) {
now := time.Now()
future := now.Add(10 * time.Minute).Format(time.RFC3339)
account := &Account{
Platform: PlatformAntigravity,
Extra: map[string]any{
modelRateLimitsKey: map[string]any{
"claude-sonnet-4-5-thinking": map[string]any{
"rate_limit_reset_at": future,
},
},
},
}
ctx := context.WithValue(context.Background(), ctxkey.ThinkingEnabled, true)
if !account.isModelRateLimitedWithContext(ctx, "claude-sonnet-4-5") {
t.Errorf("expected model to be rate limited")
}
}
func TestGetModelRateLimitRemainingTime(t *testing.T) {
now := time.Now()
future10m := now.Add(10 * time.Minute).Format(time.RFC3339)
future5m := now.Add(5 * time.Minute).Format(time.RFC3339)
past := now.Add(-10 * time.Minute).Format(time.RFC3339)
tests := []struct {
name string
account *Account
requestedModel string
minExpected time.Duration
maxExpected time.Duration
}{
{
name: "nil account",
account: nil,
requestedModel: "claude-sonnet-4-5",
minExpected: 0,
maxExpected: 0,
},
{
name: "model rate limited - direct hit",
account: &Account{
Extra: map[string]any{
modelRateLimitsKey: map[string]any{
"claude-sonnet-4-5": map[string]any{
"rate_limit_reset_at": future10m,
},
},
},
},
requestedModel: "claude-sonnet-4-5",
minExpected: 9 * time.Minute,
maxExpected: 11 * time.Minute,
},
{
name: "model rate limited - via mapping",
account: &Account{
Credentials: map[string]any{
"model_mapping": map[string]any{
"claude-3-5-sonnet": "claude-sonnet-4-5",
},
},
Extra: map[string]any{
modelRateLimitsKey: map[string]any{
"claude-sonnet-4-5": map[string]any{
"rate_limit_reset_at": future5m,
},
},
},
},
requestedModel: "claude-3-5-sonnet",
minExpected: 4 * time.Minute,
maxExpected: 6 * time.Minute,
},
{
name: "expired rate limit",
account: &Account{
Extra: map[string]any{
modelRateLimitsKey: map[string]any{
"claude-sonnet-4-5": map[string]any{
"rate_limit_reset_at": past,
},
},
},
},
requestedModel: "claude-sonnet-4-5",
minExpected: 0,
maxExpected: 0,
},
{
name: "no rate limit data",
account: &Account{},
requestedModel: "claude-sonnet-4-5",
minExpected: 0,
maxExpected: 0,
},
{
name: "no scope fallback",
account: &Account{
Extra: map[string]any{
modelRateLimitsKey: map[string]any{
"claude_sonnet": map[string]any{
"rate_limit_reset_at": future5m,
},
},
},
},
requestedModel: "claude-3-5-sonnet-20241022",
minExpected: 0,
maxExpected: 0,
},
{
name: "antigravity platform - claude-opus-4-5-thinking mapped to opus-4-6-thinking",
account: &Account{
Platform: PlatformAntigravity,
Extra: map[string]any{
modelRateLimitsKey: map[string]any{
"claude-opus-4-6-thinking": map[string]any{
"rate_limit_reset_at": future5m,
},
},
},
},
requestedModel: "claude-opus-4-5-thinking",
minExpected: 4 * time.Minute,
maxExpected: 6 * time.Minute,
},
{
name: "antigravity platform - gemini family rate limit remaining",
account: &Account{
Platform: PlatformAntigravity,
Extra: map[string]any{
modelRateLimitsKey: map[string]any{
antigravityGeminiModelRateLimitKey: map[string]any{
"rate_limit_reset_at": future10m,
},
},
},
},
requestedModel: "gemini-3-pro-preview",
minExpected: 9 * time.Minute,
maxExpected: 11 * time.Minute,
},
{
name: "antigravity platform - gemini family remaining ignored for claude",
account: &Account{
Platform: PlatformAntigravity,
Extra: map[string]any{
modelRateLimitsKey: map[string]any{
antigravityGeminiModelRateLimitKey: map[string]any{
"rate_limit_reset_at": future10m,
},
},
},
},
requestedModel: "claude-sonnet-4-5",
minExpected: 0,
maxExpected: 0,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
result := tt.account.GetModelRateLimitRemainingTimeWithContext(context.Background(), tt.requestedModel)
if result < tt.minExpected || result > tt.maxExpected {
t.Errorf("GetModelRateLimitRemainingTime() = %v, want between %v and %v", result, tt.minExpected, tt.maxExpected)
}
})
}
}
func TestGetRateLimitRemainingTime(t *testing.T) {
now := time.Now()
future15m := now.Add(15 * time.Minute).Format(time.RFC3339)
future5m := now.Add(5 * time.Minute).Format(time.RFC3339)
tests := []struct {
name string
account *Account
requestedModel string
minExpected time.Duration
maxExpected time.Duration
}{
{
name: "nil account",
account: nil,
requestedModel: "claude-sonnet-4-5",
minExpected: 0,
maxExpected: 0,
},
{
name: "model rate limited - 15 minutes",
account: &Account{
Platform: PlatformAntigravity,
Extra: map[string]any{
modelRateLimitsKey: map[string]any{
"claude-sonnet-4-5": map[string]any{
"rate_limit_reset_at": future15m,
},
},
},
},
requestedModel: "claude-sonnet-4-5",
minExpected: 14 * time.Minute,
maxExpected: 16 * time.Minute,
},
{
name: "only model rate limited",
account: &Account{
Platform: PlatformAntigravity,
Extra: map[string]any{
modelRateLimitsKey: map[string]any{
"claude-sonnet-4-5": map[string]any{
"rate_limit_reset_at": future5m,
},
},
},
},
requestedModel: "claude-sonnet-4-5",
minExpected: 4 * time.Minute,
maxExpected: 6 * time.Minute,
},
{
name: "neither rate limited",
account: &Account{
Platform: PlatformAntigravity,
},
requestedModel: "claude-sonnet-4-5",
minExpected: 0,
maxExpected: 0,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
result := tt.account.GetRateLimitRemainingTimeWithContext(context.Background(), tt.requestedModel)
if result < tt.minExpected || result > tt.maxExpected {
t.Errorf("GetRateLimitRemainingTime() = %v, want between %v and %v", result, tt.minExpected, tt.maxExpected)
}
})
}
}
func TestIsModelRateLimited_AnthropicFableFamilyKey(t *testing.T) {
now := time.Now()
future := now.Add(48 * time.Hour).Format(time.RFC3339)
account := &Account{
Platform: PlatformAnthropic,
Extra: map[string]any{
modelRateLimitsKey: map[string]any{
anthropicFableRateLimitKey: map[string]any{
"rate_limit_reset_at": future,
},
},
},
}
tests := []struct {
requestedModel string
expected bool
}{
{"claude-fable-5", true},
{"claude-fable-5[1m]", true}, // 家族 key 覆盖变体
{"Claude-Fable-5-20260601", true}, // 大小写不敏感
{"claude-sonnet-4-6", false}, // 其他模型不受影响
{"claude-opus-4-8", false},
}
for _, tc := range tests {
t.Run(tc.requestedModel, func(t *testing.T) {
got := account.isModelRateLimitedWithContext(context.Background(), tc.requestedModel)
require.Equal(t, tc.expected, got)
remaining := account.GetModelRateLimitRemainingTimeWithContext(context.Background(), tc.requestedModel)
require.Equal(t, tc.expected, remaining > 0)
})
}
}
func TestIsAnthropicFableModel(t *testing.T) {
require.True(t, isAnthropicFableModel("claude-fable-5"))
require.True(t, isAnthropicFableModel("claude-fable-5[1m]"))
require.True(t, isAnthropicFableModel("Claude-Fable-5"))
require.False(t, isAnthropicFableModel("claude-sonnet-4-6"))
require.False(t, isAnthropicFableModel(""))
}