Files
sub2api/backend/internal/service/temp_unsched_test.go
erio 5e98445b22 feat(antigravity): comprehensive enhancements - model mapping, rate limiting, scheduling & ops
Key changes:
- Upgrade model mapping: Opus 4.5 → Opus 4.6-thinking with precise matching
- Unified rate limiting: scope-level → model-level with Redis snapshot sync
- Load-balanced scheduling by call count with smart retry mechanism
- Force cache billing support
- Model identity injection in prompts with leak prevention
- Thinking mode auto-handling (max_tokens/budget_tokens fix)
- Frontend: whitelist mode toggle, model mapping validation, status indicators
- Gemini session fallback with Redis Trie O(L) matching
- Ops: enhanced concurrency monitoring, account availability, retry logic
- Migration scripts: 049-051 for model mapping unification
2026-02-07 12:31:10 +08:00

379 lines
8.4 KiB
Go
Raw Permalink Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

//go:build unit
package service
import (
"testing"
"time"
"github.com/stretchr/testify/require"
)
// ============ 临时限流单元测试 ============
// TestMatchTempUnschedKeyword 测试关键词匹配函数
func TestMatchTempUnschedKeyword(t *testing.T) {
tests := []struct {
name string
body string
keywords []string
want string
}{
{
name: "match_first",
body: "server is overloaded",
keywords: []string{"overloaded", "capacity"},
want: "overloaded",
},
{
name: "match_second",
body: "no capacity available",
keywords: []string{"overloaded", "capacity"},
want: "capacity",
},
{
name: "no_match",
body: "internal error",
keywords: []string{"overloaded", "capacity"},
want: "",
},
{
name: "empty_body",
body: "",
keywords: []string{"overloaded"},
want: "",
},
{
name: "empty_keywords",
body: "server is overloaded",
keywords: []string{},
want: "",
},
{
name: "whitespace_keyword",
body: "server is overloaded",
keywords: []string{" ", "overloaded"},
want: "overloaded",
},
{
// matchTempUnschedKeyword 期望 body 已经是小写的
// 所以要测试大小写不敏感匹配,需要传入小写的 body
name: "case_insensitive_body_lowered",
body: "server is overloaded", // body 已经是小写
keywords: []string{"OVERLOADED"}, // keyword 会被转为小写比较
want: "OVERLOADED",
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
got := matchTempUnschedKeyword(tt.body, tt.keywords)
require.Equal(t, tt.want, got)
})
}
}
// TestAccountIsSchedulable_TempUnschedulable 测试临时限流账号不可调度
func TestAccountIsSchedulable_TempUnschedulable(t *testing.T) {
future := time.Now().Add(10 * time.Minute)
past := time.Now().Add(-10 * time.Minute)
tests := []struct {
name string
account *Account
want bool
}{
{
name: "temp_unschedulable_active",
account: &Account{
Status: StatusActive,
Schedulable: true,
TempUnschedulableUntil: &future,
},
want: false,
},
{
name: "temp_unschedulable_expired",
account: &Account{
Status: StatusActive,
Schedulable: true,
TempUnschedulableUntil: &past,
},
want: true,
},
{
name: "no_temp_unschedulable",
account: &Account{
Status: StatusActive,
Schedulable: true,
TempUnschedulableUntil: nil,
},
want: true,
},
{
name: "temp_unschedulable_with_rate_limit",
account: &Account{
Status: StatusActive,
Schedulable: true,
TempUnschedulableUntil: &future,
RateLimitResetAt: &past, // 过期的限流不影响
},
want: false, // 临时限流生效
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
got := tt.account.IsSchedulable()
require.Equal(t, tt.want, got)
})
}
}
// TestAccount_IsTempUnschedulableEnabled 测试临时限流开关
func TestAccount_IsTempUnschedulableEnabled(t *testing.T) {
tests := []struct {
name string
account *Account
want bool
}{
{
name: "enabled",
account: &Account{
Credentials: map[string]any{
"temp_unschedulable_enabled": true,
},
},
want: true,
},
{
name: "disabled",
account: &Account{
Credentials: map[string]any{
"temp_unschedulable_enabled": false,
},
},
want: false,
},
{
name: "not_set",
account: &Account{
Credentials: map[string]any{},
},
want: false,
},
{
name: "nil_credentials",
account: &Account{},
want: false,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
got := tt.account.IsTempUnschedulableEnabled()
require.Equal(t, tt.want, got)
})
}
}
// TestAccount_GetTempUnschedulableRules 测试获取临时限流规则
func TestAccount_GetTempUnschedulableRules(t *testing.T) {
tests := []struct {
name string
account *Account
wantCount int
}{
{
name: "has_rules",
account: &Account{
Credentials: map[string]any{
"temp_unschedulable_rules": []any{
map[string]any{
"error_code": float64(503),
"keywords": []any{"overloaded"},
"duration_minutes": float64(5),
},
map[string]any{
"error_code": float64(500),
"keywords": []any{"internal"},
"duration_minutes": float64(10),
},
},
},
},
wantCount: 2,
},
{
name: "empty_rules",
account: &Account{
Credentials: map[string]any{
"temp_unschedulable_rules": []any{},
},
},
wantCount: 0,
},
{
name: "no_rules",
account: &Account{
Credentials: map[string]any{},
},
wantCount: 0,
},
{
name: "nil_credentials",
account: &Account{},
wantCount: 0,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
rules := tt.account.GetTempUnschedulableRules()
require.Len(t, rules, tt.wantCount)
})
}
}
// TestTempUnschedulableRule_Parse 测试规则解析
func TestTempUnschedulableRule_Parse(t *testing.T) {
account := &Account{
Credentials: map[string]any{
"temp_unschedulable_rules": []any{
map[string]any{
"error_code": float64(503),
"keywords": []any{"overloaded", "capacity"},
"duration_minutes": float64(5),
},
},
},
}
rules := account.GetTempUnschedulableRules()
require.Len(t, rules, 1)
rule := rules[0]
require.Equal(t, 503, rule.ErrorCode)
require.Equal(t, []string{"overloaded", "capacity"}, rule.Keywords)
require.Equal(t, 5, rule.DurationMinutes)
}
// TestTruncateTempUnschedMessage 测试消息截断
func TestTruncateTempUnschedMessage(t *testing.T) {
tests := []struct {
name string
body []byte
maxBytes int
want string
}{
{
name: "short_message",
body: []byte("short"),
maxBytes: 100,
want: "short",
},
{
// 截断后会 TrimSpace所以末尾的空格会被移除
name: "truncate_long_message",
body: []byte("this is a very long message that needs to be truncated"),
maxBytes: 20,
want: "this is a very long", // 截断后 TrimSpace
},
{
name: "empty_body",
body: []byte{},
maxBytes: 100,
want: "",
},
{
name: "zero_max_bytes",
body: []byte("test"),
maxBytes: 0,
want: "",
},
{
name: "whitespace_trimmed",
body: []byte(" test "),
maxBytes: 100,
want: "test",
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
got := truncateTempUnschedMessage(tt.body, tt.maxBytes)
require.Equal(t, tt.want, got)
})
}
}
// TestTempUnschedState 测试临时限流状态结构
func TestTempUnschedState(t *testing.T) {
now := time.Now()
until := now.Add(5 * time.Minute)
state := &TempUnschedState{
UntilUnix: until.Unix(),
TriggeredAtUnix: now.Unix(),
StatusCode: 503,
MatchedKeyword: "overloaded",
RuleIndex: 0,
ErrorMessage: "Server is overloaded",
}
require.Equal(t, 503, state.StatusCode)
require.Equal(t, "overloaded", state.MatchedKeyword)
require.Equal(t, 0, state.RuleIndex)
// 验证时间戳
require.Equal(t, until.Unix(), state.UntilUnix)
require.Equal(t, now.Unix(), state.TriggeredAtUnix)
}
// TestAccount_TempUnschedulableUntil 测试临时限流时间字段
func TestAccount_TempUnschedulableUntil(t *testing.T) {
future := time.Now().Add(10 * time.Minute)
past := time.Now().Add(-10 * time.Minute)
tests := []struct {
name string
account *Account
schedulable bool
}{
{
name: "active_temp_unsched_not_schedulable",
account: &Account{
Status: StatusActive,
Schedulable: true,
TempUnschedulableUntil: &future,
},
schedulable: false,
},
{
name: "expired_temp_unsched_is_schedulable",
account: &Account{
Status: StatusActive,
Schedulable: true,
TempUnschedulableUntil: &past,
},
schedulable: true,
},
{
name: "nil_temp_unsched_is_schedulable",
account: &Account{
Status: StatusActive,
Schedulable: true,
TempUnschedulableUntil: nil,
},
schedulable: true,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
got := tt.account.IsSchedulable()
require.Equal(t, tt.schedulable, got)
})
}
}