test: add warmup request interception unit tests
Add comprehensive tests for warmup request interception behavior covering Antigravity accounts with various credential configurations.
This commit is contained in:
erio
@@ -0,0 +1,340 @@
|
||||
//go:build unit
|
||||
|
||||
package handler
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/ctxkey"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/pagination"
|
||||
middleware "github.com/Wei-Shaw/sub2api/internal/server/middleware"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
// 目标：严格验证“antigravity 账号通过 /v1/messages 提供 Claude 服务时”，
|
||||
// 当账号 credentials.intercept_warmup_requests=true 且请求为 Warmup 时，
|
||||
// 后端会在转发上游前直接拦截并返回 mock 响应（不依赖上游）。
|
||||
|
||||
type fakeSchedulerCache struct {
|
||||
accounts []*service.Account
|
||||
}
|
||||
|
||||
func (f *fakeSchedulerCache) GetSnapshot(_ context.Context, _ service.SchedulerBucket) ([]*service.Account, bool, error) {
|
||||
return f.accounts, true, nil
|
||||
}
|
||||
func (f *fakeSchedulerCache) SetSnapshot(_ context.Context, _ service.SchedulerBucket, _ []service.Account) error {
|
||||
return nil
|
||||
}
|
||||
func (f *fakeSchedulerCache) GetAccount(_ context.Context, _ int64) (*service.Account, error) {
|
||||
return nil, nil
|
||||
}
|
||||
func (f *fakeSchedulerCache) SetAccount(_ context.Context, _ *service.Account) error { return nil }
|
||||
func (f *fakeSchedulerCache) DeleteAccount(_ context.Context, _ int64) error { return nil }
|
||||
func (f *fakeSchedulerCache) UpdateLastUsed(_ context.Context, _ map[int64]time.Time) error {
|
||||
return nil
|
||||
}
|
||||
func (f *fakeSchedulerCache) TryLockBucket(_ context.Context, _ service.SchedulerBucket, _ time.Duration) (bool, error) {
|
||||
return true, nil
|
||||
}
|
||||
func (f *fakeSchedulerCache) ListBuckets(_ context.Context) ([]service.SchedulerBucket, error) {
|
||||
return nil, nil
|
||||
}
|
||||
func (f *fakeSchedulerCache) GetOutboxWatermark(_ context.Context) (int64, error) { return 0, nil }
|
||||
func (f *fakeSchedulerCache) SetOutboxWatermark(_ context.Context, _ int64) error { return nil }
|
||||
|
||||
type fakeGroupRepo struct {
|
||||
group *service.Group
|
||||
}
|
||||
|
||||
func (f *fakeGroupRepo) Create(context.Context, *service.Group) error { return nil }
|
||||
func (f *fakeGroupRepo) GetByID(context.Context, int64) (*service.Group, error) {
|
||||
return f.group, nil
|
||||
}
|
||||
func (f *fakeGroupRepo) GetByIDLite(context.Context, int64) (*service.Group, error) {
|
||||
return f.group, nil
|
||||
}
|
||||
func (f *fakeGroupRepo) Update(context.Context, *service.Group) error { return nil }
|
||||
func (f *fakeGroupRepo) Delete(context.Context, int64) error { return nil }
|
||||
func (f *fakeGroupRepo) DeleteCascade(context.Context, int64) ([]int64, error) { return nil, nil }
|
||||
func (f *fakeGroupRepo) List(context.Context, pagination.PaginationParams) ([]service.Group, *pagination.PaginationResult, error) {
|
||||
return nil, nil, nil
|
||||
}
|
||||
func (f *fakeGroupRepo) ListWithFilters(context.Context, pagination.PaginationParams, string, string, string, *bool) ([]service.Group, *pagination.PaginationResult, error) {
|
||||
return nil, nil, nil
|
||||
}
|
||||
func (f *fakeGroupRepo) ListActive(context.Context) ([]service.Group, error) { return nil, nil }
|
||||
func (f *fakeGroupRepo) ListActiveByPlatform(context.Context, string) ([]service.Group, error) {
|
||||
return nil, nil
|
||||
}
|
||||
func (f *fakeGroupRepo) ExistsByName(context.Context, string) (bool, error) { return false, nil }
|
||||
func (f *fakeGroupRepo) GetAccountCount(context.Context, int64) (int64, error) { return 0, nil }
|
||||
func (f *fakeGroupRepo) DeleteAccountGroupsByGroupID(context.Context, int64) (int64, error) {
|
||||
return 0, nil
|
||||
}
|
||||
func (f *fakeGroupRepo) GetAccountIDsByGroupIDs(context.Context, []int64) ([]int64, error) {
|
||||
return nil, nil
|
||||
}
|
||||
func (f *fakeGroupRepo) BindAccountsToGroup(context.Context, int64, []int64) error { return nil }
|
||||
func (f *fakeGroupRepo) UpdateSortOrders(context.Context, []service.GroupSortOrderUpdate) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
type fakeConcurrencyCache struct{}
|
||||
|
||||
func (f *fakeConcurrencyCache) AcquireAccountSlot(context.Context, int64, int, string) (bool, error) {
|
||||
return true, nil
|
||||
}
|
||||
func (f *fakeConcurrencyCache) ReleaseAccountSlot(context.Context, int64, string) error { return nil }
|
||||
func (f *fakeConcurrencyCache) GetAccountConcurrency(context.Context, int64) (int, error) {
|
||||
return 0, nil
|
||||
}
|
||||
func (f *fakeConcurrencyCache) IncrementAccountWaitCount(context.Context, int64, int) (bool, error) {
|
||||
return true, nil
|
||||
}
|
||||
func (f *fakeConcurrencyCache) DecrementAccountWaitCount(context.Context, int64) error { return nil }
|
||||
func (f *fakeConcurrencyCache) GetAccountWaitingCount(context.Context, int64) (int, error) {
|
||||
return 0, nil
|
||||
}
|
||||
func (f *fakeConcurrencyCache) AcquireUserSlot(context.Context, int64, int, string) (bool, error) {
|
||||
return true, nil
|
||||
}
|
||||
func (f *fakeConcurrencyCache) ReleaseUserSlot(context.Context, int64, string) error { return nil }
|
||||
func (f *fakeConcurrencyCache) GetUserConcurrency(context.Context, int64) (int, error) { return 0, nil }
|
||||
func (f *fakeConcurrencyCache) IncrementWaitCount(context.Context, int64, int) (bool, error) {
|
||||
return true, nil
|
||||
}
|
||||
func (f *fakeConcurrencyCache) DecrementWaitCount(context.Context, int64) error { return nil }
|
||||
func (f *fakeConcurrencyCache) GetAccountsLoadBatch(context.Context, []service.AccountWithConcurrency) (map[int64]*service.AccountLoadInfo, error) {
|
||||
return map[int64]*service.AccountLoadInfo{}, nil
|
||||
}
|
||||
func (f *fakeConcurrencyCache) GetUsersLoadBatch(context.Context, []service.UserWithConcurrency) (map[int64]*service.UserLoadInfo, error) {
|
||||
return map[int64]*service.UserLoadInfo{}, nil
|
||||
}
|
||||
func (f *fakeConcurrencyCache) CleanupExpiredAccountSlots(context.Context, int64) error { return nil }
|
||||
|
||||
func newTestGatewayHandler(t *testing.T, group *service.Group, accounts []*service.Account) (*GatewayHandler, func()) {
|
||||
t.Helper()
|
||||
|
||||
schedulerCache := &fakeSchedulerCache{accounts: accounts}
|
||||
schedulerSnapshot := service.NewSchedulerSnapshotService(schedulerCache, nil, nil, nil, nil)
|
||||
|
||||
gwSvc := service.NewGatewayService(
|
||||
nil, // accountRepo (not used: scheduler snapshot hit)
|
||||
&fakeGroupRepo{group: group},
|
||||
nil, // usageLogRepo
|
||||
nil, // userRepo
|
||||
nil, // userSubRepo
|
||||
nil, // userGroupRateRepo
|
||||
nil, // cache (disable sticky)
|
||||
nil, // cfg
|
||||
schedulerSnapshot,
|
||||
nil, // concurrencyService (disable load-aware; tryAcquire always acquired)
|
||||
nil, // billingService
|
||||
nil, // rateLimitService
|
||||
nil, // billingCacheService
|
||||
nil, // identityService
|
||||
nil, // httpUpstream
|
||||
nil, // deferredService
|
||||
nil, // claudeTokenProvider
|
||||
nil, // sessionLimitCache
|
||||
nil, // digestStore
|
||||
)
|
||||
|
||||
// RunModeSimple：跳过计费检查，避免引入 repo/cache 依赖。
|
||||
cfg := &config.Config{RunMode: config.RunModeSimple}
|
||||
billingCacheSvc := service.NewBillingCacheService(nil, nil, nil, cfg)
|
||||
|
||||
concurrencySvc := service.NewConcurrencyService(&fakeConcurrencyCache{})
|
||||
concurrencyHelper := NewConcurrencyHelper(concurrencySvc, SSEPingFormatClaude, 0)
|
||||
|
||||
h := &GatewayHandler{
|
||||
gatewayService: gwSvc,
|
||||
billingCacheService: billingCacheSvc,
|
||||
concurrencyHelper: concurrencyHelper,
|
||||
// 这些字段对本测试不敏感，保持较小即可
|
||||
maxAccountSwitches: 1,
|
||||
maxAccountSwitchesGemini: 1,
|
||||
}
|
||||
|
||||
cleanup := func() {
|
||||
billingCacheSvc.Stop()
|
||||
}
|
||||
return h, cleanup
|
||||
}
|
||||
|
||||
func TestGatewayHandlerMessages_InterceptWarmup_AntigravityAccount_MixedSchedulingV1(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
groupID := int64(2001)
|
||||
accountID := int64(1001)
|
||||
|
||||
group := &service.Group{
|
||||
ID: groupID,
|
||||
Hydrated: true,
|
||||
Platform: service.PlatformAnthropic, // /v1/messages（Claude兼容）入口
|
||||
Status: service.StatusActive,
|
||||
}
|
||||
|
||||
account := &service.Account{
|
||||
ID: accountID,
|
||||
Name: "ag-1",
|
||||
Platform: service.PlatformAntigravity,
|
||||
Type: service.AccountTypeOAuth,
|
||||
Credentials: map[string]any{
|
||||
"access_token": "tok_xxx",
|
||||
"intercept_warmup_requests": true,
|
||||
},
|
||||
Extra: map[string]any{
|
||||
"mixed_scheduling": true, // 关键：允许被 anthropic 分组混合调度选中
|
||||
},
|
||||
Concurrency: 1,
|
||||
Priority: 1,
|
||||
Status: service.StatusActive,
|
||||
Schedulable: true,
|
||||
AccountGroups: []service.AccountGroup{{AccountID: accountID, GroupID: groupID}},
|
||||
}
|
||||
|
||||
h, cleanup := newTestGatewayHandler(t, group, []*service.Account{account})
|
||||
defer cleanup()
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(rec)
|
||||
|
||||
body := []byte(`{
|
||||
"model": "claude-sonnet-4-5",
|
||||
"max_tokens": 256,
|
||||
"messages": [{"role":"user","content":[{"type":"text","text":"Warmup"}]}]
|
||||
}`)
|
||||
req := httptest.NewRequest("POST", "/v1/messages", bytes.NewReader(body))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
req = req.WithContext(context.WithValue(req.Context(), ctxkey.Group, group))
|
||||
c.Request = req
|
||||
|
||||
apiKey := &service.APIKey{
|
||||
ID: 3001,
|
||||
UserID: 4001,
|
||||
GroupID: &groupID,
|
||||
Status: service.StatusActive,
|
||||
User: &service.User{
|
||||
ID: 4001,
|
||||
Concurrency: 10,
|
||||
Balance: 100,
|
||||
},
|
||||
Group: group,
|
||||
}
|
||||
|
||||
c.Set(string(middleware.ContextKeyAPIKey), apiKey)
|
||||
c.Set(string(middleware.ContextKeyUser), middleware.AuthSubject{UserID: apiKey.UserID, Concurrency: 10})
|
||||
|
||||
h.Messages(c)
|
||||
|
||||
require.Equal(t, 200, rec.Code)
|
||||
|
||||
// 断言：确实选中了 antigravity 账号（不是纯函数测试，而是从 Handler 里验证调度结果）
|
||||
selected, ok := c.Get(opsAccountIDKey)
|
||||
require.True(t, ok)
|
||||
require.Equal(t, accountID, selected)
|
||||
|
||||
var resp map[string]any
|
||||
require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &resp))
|
||||
require.Equal(t, "msg_mock_warmup", resp["id"])
|
||||
require.Equal(t, "claude-sonnet-4-5", resp["model"])
|
||||
|
||||
content, ok := resp["content"].([]any)
|
||||
require.True(t, ok)
|
||||
require.Len(t, content, 1)
|
||||
first, ok := content[0].(map[string]any)
|
||||
require.True(t, ok)
|
||||
require.Equal(t, "New Conversation", first["text"])
|
||||
}
|
||||
|
||||
func TestGatewayHandlerMessages_InterceptWarmup_AntigravityAccount_ForcePlatform(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
groupID := int64(2002)
|
||||
accountID := int64(1002)
|
||||
|
||||
group := &service.Group{
|
||||
ID: groupID,
|
||||
Hydrated: true,
|
||||
Platform: service.PlatformAntigravity,
|
||||
Status: service.StatusActive,
|
||||
}
|
||||
|
||||
account := &service.Account{
|
||||
ID: accountID,
|
||||
Name: "ag-2",
|
||||
Platform: service.PlatformAntigravity,
|
||||
Type: service.AccountTypeOAuth,
|
||||
Credentials: map[string]any{
|
||||
"access_token": "tok_xxx",
|
||||
"intercept_warmup_requests": true,
|
||||
},
|
||||
Concurrency: 1,
|
||||
Priority: 1,
|
||||
Status: service.StatusActive,
|
||||
Schedulable: true,
|
||||
AccountGroups: []service.AccountGroup{{AccountID: accountID, GroupID: groupID}},
|
||||
}
|
||||
|
||||
h, cleanup := newTestGatewayHandler(t, group, []*service.Account{account})
|
||||
defer cleanup()
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(rec)
|
||||
|
||||
body := []byte(`{
|
||||
"model": "claude-sonnet-4-5",
|
||||
"max_tokens": 256,
|
||||
"messages": [{"role":"user","content":[{"type":"text","text":"Warmup"}]}]
|
||||
}`)
|
||||
req := httptest.NewRequest("POST", "/antigravity/v1/messages", bytes.NewReader(body))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
|
||||
// 模拟 routes/gateway.go 里的 ForcePlatform 中间件效果：
|
||||
// - 写入 request.Context（Service读取）
|
||||
// - 写入 gin.Context（Handler快速读取）
|
||||
ctx := context.WithValue(req.Context(), ctxkey.Group, group)
|
||||
ctx = context.WithValue(ctx, ctxkey.ForcePlatform, service.PlatformAntigravity)
|
||||
req = req.WithContext(ctx)
|
||||
c.Request = req
|
||||
c.Set(string(middleware.ContextKeyForcePlatform), service.PlatformAntigravity)
|
||||
|
||||
apiKey := &service.APIKey{
|
||||
ID: 3002,
|
||||
UserID: 4002,
|
||||
GroupID: &groupID,
|
||||
Status: service.StatusActive,
|
||||
User: &service.User{
|
||||
ID: 4002,
|
||||
Concurrency: 10,
|
||||
Balance: 100,
|
||||
},
|
||||
Group: group,
|
||||
}
|
||||
|
||||
c.Set(string(middleware.ContextKeyAPIKey), apiKey)
|
||||
c.Set(string(middleware.ContextKeyUser), middleware.AuthSubject{UserID: apiKey.UserID, Concurrency: 10})
|
||||
|
||||
h.Messages(c)
|
||||
|
||||
require.Equal(t, 200, rec.Code)
|
||||
|
||||
selected, ok := c.Get(opsAccountIDKey)
|
||||
require.True(t, ok)
|
||||
require.Equal(t, accountID, selected)
|
||||
|
||||
var resp map[string]any
|
||||
require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &resp))
|
||||
require.Equal(t, "msg_mock_warmup", resp["id"])
|
||||
require.Equal(t, "claude-sonnet-4-5", resp["model"])
|
||||
}
|
||||
Reference in New Issue
Block a user