//go:build unit package service import ( "context" "errors" "testing" "github.com/Wei-Shaw/sub2api/internal/config" "github.com/Wei-Shaw/sub2api/internal/pkg/ctxkey" "github.com/stretchr/testify/require" ) func gatewayProfitTestGroup(id int64, platform string) *Group { return &Group{ ID: id, Name: "profit-" + platform, Platform: platform, Status: StatusActive, Hydrated: true, RateMultiplier: 0.5, SubscriptionType: SubscriptionTypeStandard, ProfitControlEnabled: true, ProfitMinMargin: 0, ProfitSafetyBuffer: 0, } } func gatewayProfitTestContext(group *Group) context.Context { ctx := context.WithValue(context.Background(), ctxkey.Group, group) ctx, _ = WithGatewayTokenRequestPricing(ctx) return ctx } func gatewayProfitTestAccount(id int64, platform string, rate float64, groupID int64) Account { return Account{ ID: id, Name: "account", Platform: platform, Type: AccountTypeAPIKey, Status: StatusActive, Schedulable: true, Concurrency: 2, Priority: 1, RateMultiplier: &rate, AccountGroups: []AccountGroup{{AccountID: id, GroupID: groupID}}, GroupIDs: []int64{groupID}, } } func TestGatewayProfitControlInstallsForFivePlatformsOnlyOnTokenRequests(t *testing.T) { for _, platform := range []string{ PlatformOpenAI, PlatformAnthropic, PlatformGemini, PlatformGrok, PlatformAntigravity, } { t.Run(platform, func(t *testing.T) { group := gatewayProfitTestGroup(101, platform) groupID := group.ID svc := &GatewayService{} tokenCtx := svc.withGatewayProfitControlGate(gatewayProfitTestContext(group), &groupID) gate, _ := tokenCtx.Value(openAIProfitControlGateCtxKey{}).(*openAIProfitControlGate) require.NotNil(t, gate) require.Equal(t, platform, gate.platform) require.InDelta(t, 0.5, gate.threshold, 1e-12) metadataCtx := context.WithValue(context.Background(), ctxkey.Group, group) metadataCtx = svc.withGatewayProfitControlGate(metadataCtx, &groupID) gate, _ = metadataCtx.Value(openAIProfitControlGateCtxKey{}).(*openAIProfitControlGate) require.Nil(t, gate, "未显式标记为 token 请求的入口不得装门") }) } } func TestGatewayProfitControlCompositeBillingUsesScheduledMemberConfig(t *testing.T) { billingGroup := &Group{ ID: 201, Platform: PlatformComposite, Status: StatusActive, Hydrated: true, RateMultiplier: 0.4, SubscriptionType: SubscriptionTypeStandard, } memberGroup := gatewayProfitTestGroup(202, PlatformAnthropic) memberGroup.RateMultiplier = 99 memberGroup.ProfitMinMargin = 0.25 ctx := context.WithValue(context.Background(), ctxkey.Group, billingGroup) ctx, pricingAt := WithGatewayTokenRequestPricing(ctx) svc := &GatewayService{ schedulerSnapshot: NewSchedulerSnapshotService( nil, nil, nil, profitControlGroupRepo{group: memberGroup}, nil, ), } ctx = svc.withGatewayProfitControlGate(ctx, &memberGroup.ID) gate, _ := ctx.Value(openAIProfitControlGateCtxKey{}).(*openAIProfitControlGate) require.NotNil(t, gate) require.Equal(t, memberGroup.ID, gate.groupID) require.Equal(t, PlatformAnthropic, gate.platform) require.Equal(t, pricingAt, gate.pricingAt) require.InDelta(t, 0.4*(1-0.25), gate.threshold, 1e-12, "D 必须取 composite 计费父分组,margin 取被调度成员分组") } func TestGatewayProfitControlGroupLoadFailureClearsForeignGate(t *testing.T) { billingGroup := &Group{ ID: 211, Platform: PlatformComposite, Status: StatusActive, Hydrated: true, RateMultiplier: 0.4, SubscriptionType: SubscriptionTypeStandard, } targetGroupID := int64(212) ctx := gatewayProfitTestContext(billingGroup) ctx = context.WithValue(ctx, openAIProfitControlGateCtxKey{}, &openAIProfitControlGate{ groupID: 210, platform: PlatformAnthropic, threshold: 0.1, }) svc := &GatewayService{ schedulerSnapshot: NewSchedulerSnapshotService( nil, nil, nil, profitControlFailingGroupRepo{}, nil, ), } ctx = svc.withGatewayProfitControlGate(ctx, &targetGroupID) gate, ok := ctx.Value(openAIProfitControlGateCtxKey{}).(*openAIProfitControlGate) require.True(t, ok) require.Nil(t, gate, "加载新分组失败时必须清除其他分组遗留的门") account := gatewayProfitTestAccount(213, PlatformAnthropic, 0.8, targetGroupID) require.True(t, svc.isGatewayAccountProfitEligible(ctx, &account), "配置读取失败按既定语义 fail-open") } type profitControlFailingGroupRepo struct { GroupRepository } func (profitControlFailingGroupRepo) GetByIDLite(context.Context, int64) (*Group, error) { return nil, errors.New("group cache unavailable") } // 见 profitControlGroupRepo.GetByID:利润门必须走不带账号计数聚合的 lite 读取。 func (profitControlFailingGroupRepo) GetByID(context.Context, int64) (*Group, error) { panic("profit control gate must read groups via GetByIDLite (no account-count aggregation)") } func TestGatewayProfitControlLegacyMixedAndRoutedSelection(t *testing.T) { t.Run("legacy single-platform selection", func(t *testing.T) { group := gatewayProfitTestGroup(111, PlatformGrok) cheap := gatewayProfitTestAccount(1, PlatformGrok, 0.2, group.ID) expensive := gatewayProfitTestAccount(2, PlatformGrok, 0.8, group.ID) repo := &mockAccountRepoForPlatform{ accounts: []Account{expensive, cheap}, accountsByID: map[int64]*Account{cheap.ID: &cheap, expensive.ID: &expensive}, } svc := &GatewayService{ accountRepo: repo, cache: &mockGatewayCacheForPlatform{}, cfg: testConfig(), } selected, err := svc.SelectAccountForModelWithExclusions( gatewayProfitTestContext(group), &group.ID, "", "", nil, ) require.NoError(t, err) require.Equal(t, cheap.ID, selected.ID) _, err = svc.SelectAccountForModelWithExclusions( gatewayProfitTestContext(group), &group.ID, "", "", map[int64]struct{}{cheap.ID: {}}, ) require.Error(t, err) require.ErrorIs(t, err, ErrNoAvailableAccounts) }) t.Run("mixed routing filters the routed account", func(t *testing.T) { group := gatewayProfitTestGroup(112, PlatformAnthropic) group.ModelRoutingEnabled = true group.ModelRouting = map[string][]int64{"claude-test": {2, 1}} cheap := gatewayProfitTestAccount(1, PlatformAntigravity, 0.2, group.ID) cheap.Extra = map[string]any{"mixed_scheduling": true} cheap.Credentials = map[string]any{"model_mapping": map[string]any{"claude-test": "claude-test"}} expensive := gatewayProfitTestAccount(2, PlatformAnthropic, 0.8, group.ID) repo := &mockAccountRepoForPlatform{ accounts: []Account{expensive, cheap}, accountsByID: map[int64]*Account{cheap.ID: &cheap, expensive.ID: &expensive}, } svc := &GatewayService{ accountRepo: repo, cache: &mockGatewayCacheForPlatform{}, cfg: testConfig(), } selected, err := svc.SelectAccountForModelWithExclusions( gatewayProfitTestContext(group), &group.ID, "", "claude-test", nil, ) require.NoError(t, err) require.Equal(t, cheap.ID, selected.ID) }) } func TestGatewayProfitControlLoadAwareSelectionAndFailover(t *testing.T) { group := gatewayProfitTestGroup(121, PlatformGrok) cheap := gatewayProfitTestAccount(1, PlatformGrok, 0.2, group.ID) expensive := gatewayProfitTestAccount(2, PlatformGrok, 0.8, group.ID) repo := &mockAccountRepoForPlatform{ accounts: []Account{expensive, cheap}, accountsByID: map[int64]*Account{cheap.ID: &cheap, expensive.ID: &expensive}, } cfg := &config.Config{RunMode: config.RunModeStandard} cfg.Gateway.Scheduling.LoadBatchEnabled = true svc := &GatewayService{ accountRepo: repo, cache: &mockGatewayCacheForPlatform{}, cfg: cfg, concurrencyService: NewConcurrencyService(stubConcurrencyCache{}), } result, err := svc.SelectAccountWithLoadAwareness( gatewayProfitTestContext(group), &group.ID, "", "", nil, "", 0, ) require.NoError(t, err) require.NotNil(t, result) require.Equal(t, cheap.ID, result.Account.ID) if result.ReleaseFunc != nil { result.ReleaseFunc() } result, err = svc.SelectAccountWithLoadAwareness( gatewayProfitTestContext(group), &group.ID, "", "", map[int64]struct{}{cheap.ID: {}}, "", 0, ) require.Nil(t, result) require.Error(t, err) require.ErrorIs(t, err, ErrNoAvailableAccounts) } func TestGatewayProfitControlStickyVetoKeepsBindingUntilRateRecovers(t *testing.T) { group := gatewayProfitTestGroup(131, PlatformAnthropic) expensive := gatewayProfitTestAccount(1, PlatformAnthropic, 0.8, group.ID) cheap := gatewayProfitTestAccount(2, PlatformAnthropic, 0.2, group.ID) repo := &mockAccountRepoForPlatform{ accounts: []Account{expensive, cheap}, accountsByID: map[int64]*Account{expensive.ID: &expensive, cheap.ID: &cheap}, } cache := &mockGatewayCacheForPlatform{ sessionBindings: map[string]int64{"sticky-profit": expensive.ID}, } svc := &GatewayService{ accountRepo: repo, cache: cache, cfg: testConfig(), } ctx := gatewayProfitTestContext(group) selected, err := svc.SelectAccountForModelWithExclusions(ctx, &group.ID, "sticky-profit", "", nil) require.NoError(t, err) require.Equal(t, cheap.ID, selected.ID) require.Equal(t, expensive.ID, cache.sessionBindings["sticky-profit"], "候选过滤不得覆盖旧粘性绑定") require.NoError(t, svc.BindStickySessionAfterProfitAdmission( svc.withGatewayProfitControlGate(ctx, &group.ID), &group.ID, "sticky-profit", cheap.ID, )) require.Equal(t, expensive.ID, cache.sessionBindings["sticky-profit"], "终检通过的 fallback 账号也不得覆盖旧绑定") require.Zero(t, cache.deletedSessions["sticky-profit"]) recovered := expensive recoveredRate := 0.2 recovered.RateMultiplier = &recoveredRate repo.accounts[0] = recovered repo.accountsByID[recovered.ID] = &repo.accounts[0] selected, err = svc.SelectAccountForModelWithExclusions(ctx, &group.ID, "sticky-profit", "", nil) require.NoError(t, err) require.Equal(t, recovered.ID, selected.ID, "倍率恢复后应重新命中原粘性账号") require.Zero(t, cache.deletedSessions["sticky-profit"]) } type gatewayProfitSnapshotCache struct { SchedulerCache account *Account err error } func (c *gatewayProfitSnapshotCache) GetAccount(context.Context, int64) (*Account, error) { return c.account, c.err } type gatewayProfitAccountRepo struct { AccountRepository account *Account err error } func (r gatewayProfitAccountRepo) GetByID(context.Context, int64) (*Account, error) { return r.account, r.err } func TestGatewayProfitControlTerminalRefreshUsesReplacementObject(t *testing.T) { selected := gatewayProfitTestAccount(141, PlatformGemini, 0.2, 1) replacement := selected expensiveRate := 0.8 replacement.RateMultiplier = &expensiveRate snapshot := NewSchedulerSnapshotService( &gatewayProfitSnapshotCache{account: &replacement}, nil, gatewayProfitAccountRepo{}, nil, nil, ) ctx := context.WithValue(context.Background(), openAIProfitControlGateCtxKey{}, &openAIProfitControlGate{ groupID: 1, platform: PlatformGemini, threshold: 0.5, }) latest, vetoed, reason := profitControlVetoLatest(ctx, &selected, snapshot) require.Same(t, &replacement, latest) require.True(t, vetoed) require.Equal(t, openAIProfitFilterReasonThreshold, reason) require.InDelta(t, 0.2, *selected.RateMultiplier, 1e-12, "测试必须替换缓存对象,不能原地修改旧指针") } func TestGatewayProfitControlTerminalRefreshFallsBackFromCacheToDatabase(t *testing.T) { selected := gatewayProfitTestAccount(145, PlatformAnthropic, 0.2, 1) replacement := selected expensiveRate := 0.8 replacement.RateMultiplier = &expensiveRate snapshot := NewSchedulerSnapshotService( &gatewayProfitSnapshotCache{err: errors.New("cache unavailable")}, nil, gatewayProfitAccountRepo{account: &replacement}, nil, nil, ) ctx := context.WithValue(context.Background(), openAIProfitControlGateCtxKey{}, &openAIProfitControlGate{ groupID: 1, platform: PlatformAnthropic, threshold: 0.5, }) latest, vetoed, reason := profitControlVetoLatest(ctx, &selected, snapshot) require.Same(t, &replacement, latest) require.True(t, vetoed, "缓存读取失败时必须继续从数据库重读,不能直接使用选号旧对象") require.Equal(t, openAIProfitFilterReasonThreshold, reason) } func TestGatewayProfitControlTerminalRefreshFailureFallsBackToSelectedObject(t *testing.T) { selected := gatewayProfitTestAccount(151, PlatformAntigravity, 0.2, 1) snapshot := NewSchedulerSnapshotService( &gatewayProfitSnapshotCache{err: errors.New("cache unavailable")}, nil, gatewayProfitAccountRepo{err: errors.New("database unavailable")}, nil, nil, ) ctx := context.WithValue(context.Background(), openAIProfitControlGateCtxKey{}, &openAIProfitControlGate{ groupID: 1, platform: PlatformAntigravity, threshold: 0.5, }) latest, vetoed, reason := profitControlVetoLatest(ctx, &selected, snapshot) require.Same(t, &selected, latest) require.False(t, vetoed) require.Empty(t, reason) } // 选号结果携带门:门安装在调度栈局部 ctx 上,handler 必须经 // ContextWithSelectionProfitGate 重放后终检与准入后绑定才可见(评审修复回归)。 func TestGatewayProfitControlSelectionCarriesGateToHandlerContext(t *testing.T) { group := gatewayProfitTestGroup(1, PlatformAnthropic) svc := &GatewayService{} expensive := gatewayProfitTestAccount(161, PlatformAnthropic, 0.9, group.ID) gateCtx := svc.withGatewayProfitControlGate(gatewayProfitTestContext(group), &group.ID) selection, err := svc.newSelectionResult(gateCtx, &expensive, true, nil, nil) require.NoError(t, err) require.True(t, selection.ProfitGateActive(), "选号结果必须携带调度栈内生效的门") // 修复前的缺陷形态:handler 原始 ctx 不含门,终检退化为空操作。 _, vetoed, _ := svc.GatewayProfitControlVetoLatest(context.Background(), &expensive) require.False(t, vetoed, "对照组:不重放门时终检确实看不到门") handlerCtx := ContextWithSelectionProfitGate(context.Background(), selection) latest, vetoed, reason := svc.GatewayProfitControlVetoLatest(handlerCtx, &expensive) require.True(t, vetoed, "重放门后终检必须真实生效") require.Equal(t, openAIProfitFilterReasonThreshold, reason) require.NotNil(t, latest) // 无门选号不携带门,重放为无操作。 plain, err := svc.newSelectionResult(context.Background(), &expensive, true, nil, nil) require.NoError(t, err) require.False(t, plain.ProfitGateActive()) require.Equal(t, context.Background(), ContextWithSelectionProfitGate(context.Background(), plain)) } // 生图意图不关门(H1/H2 回归锚点):/v1/responses 混合请求即使带生图声明, // token 定价上下文照常装配,共享门照常安装并否决越线账号。 func TestGatewayProfitControlImageIntentDoesNotDisableGate(t *testing.T) { group := gatewayProfitTestGroup(2, PlatformAnthropic) svc := &GatewayService{} expensive := gatewayProfitTestAccount(162, PlatformAnthropic, 0.9, group.ID) ctx := gatewayProfitTestContext(group) ctx = WithOpenAIImageGenerationIntent(ctx) gateCtx := svc.withGatewayProfitControlGate(ctx, &group.ID) require.False(t, svc.isGatewayAccountProfitEligible(gateCtx, &expensive), "请求体里的生图声明(含被动 image_gen namespace)不得关闭利润门") } // 无门时准入后绑定回退官方 eager 语义;门下读失败保守不写(评审 M-Bind 回归)。 func TestGatewayProfitControlAfterAdmissionBindSemantics(t *testing.T) { groupID := int64(3) expensiveID := int64(171) cheapID := int64(172) t.Run("eager without gate", func(t *testing.T) { cache := &mockGatewayCacheForPlatform{sessionBindings: map[string]int64{"s": expensiveID}} svc := &GatewayService{cache: cache} require.NoError(t, svc.BindStickySessionAfterProfitAdmission(context.Background(), &groupID, "s", cheapID)) require.Equal(t, cheapID, cache.sessionBindings["s"], "无门时保持既有 eager 绑定行为") }) t.Run("gated read failure is conservative", func(t *testing.T) { // mock 的 miss 返回非 sentinel 错误,等价于 Redis 读失败:门下保守不写。 cache := &mockGatewayCacheForPlatform{sessionBindings: map[string]int64{}} svc := &GatewayService{cache: cache} gate := &openAIProfitControlGate{groupID: groupID, platform: PlatformAnthropic, threshold: 0.5} gateCtx := context.WithValue(context.Background(), openAIProfitControlGateCtxKey{}, gate) require.NoError(t, svc.BindStickySessionAfterProfitAdmission(gateCtx, &groupID, "absent", cheapID)) require.NotContains(t, cache.sessionBindings, "absent") }) t.Run("gated sentinel miss binds", func(t *testing.T) { cache := &sentinelMissGatewayCache{mockGatewayCacheForPlatform: &mockGatewayCacheForPlatform{sessionBindings: map[string]int64{}}} svc := &GatewayService{cache: cache} gate := &openAIProfitControlGate{groupID: groupID, platform: PlatformAnthropic, threshold: 0.5} gateCtx := context.WithValue(context.Background(), openAIProfitControlGateCtxKey{}, gate) require.NoError(t, svc.BindStickySessionAfterProfitAdmission(gateCtx, &groupID, "fresh", cheapID)) require.Equal(t, cheapID, cache.sessionBindings["fresh"], "门下无既有绑定(sentinel miss)应建立粘性") }) } // sentinelMissGatewayCache 让 miss 返回与真实仓库一致的 ErrStickySessionNotFound。 type sentinelMissGatewayCache struct { *mockGatewayCacheForPlatform } func (c *sentinelMissGatewayCache) GetSessionAccountID(ctx context.Context, groupID int64, sessionHash string) (int64, error) { if id, ok := c.sessionBindings[sessionHash]; ok { return id, nil } return 0, ErrStickySessionNotFound }