5123d92b44
Adds an opt-in tier-based fallback scheduling path for Anthropic accounts: - accountTierLevel(): derives tier from account type without DB migration (tier-0=OAuth/SetupToken, tier-1=APIKey, tier-2=Bedrock) - enableTierFallbackChain(): new config flag gateway.scheduling.enable_tier_fallback_chain (default false) - selectAccountWithTierFallback(): loads all Anthropic accounts, groups by tier, honors sticky sessions, applies all existing schedulability guards, then tries tiers 0→1→2 in order via tryAcquireByLegacyOrder - Wired into SelectAccountForModelWithExclusions: Anthropic platform + tier fallback enabled → calls new path instead of mixed scheduling - Fix pre-existing unit-test build break: NewGatewayService now requires *RPMTokenBucketService (added in Task #5); add missing nil param - 7 tests: tier mapping, config toggle, subscription preference, APIKey fallback, exclusion handling, empty-pool error, Bedrock last resort
139 lines
6.2 KiB
Go
139 lines
6.2 KiB
Go
//go:build unit
|
|
|
|
package service
|
|
|
|
import (
|
|
"context"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/Wei-Shaw/sub2api/internal/config"
|
|
"github.com/stretchr/testify/require"
|
|
)
|
|
|
|
func TestAccountTierLevel(t *testing.T) {
|
|
require.Equal(t, 0, accountTierLevel(nil))
|
|
require.Equal(t, 0, accountTierLevel(&Account{Type: AccountTypeOAuth}))
|
|
require.Equal(t, 0, accountTierLevel(&Account{Type: AccountTypeSetupToken}))
|
|
require.Equal(t, 0, accountTierLevel(&Account{Type: "unknown"}))
|
|
require.Equal(t, 1, accountTierLevel(&Account{Type: AccountTypeAPIKey}))
|
|
require.Equal(t, 2, accountTierLevel(&Account{Type: AccountTypeBedrock}))
|
|
}
|
|
|
|
func TestGatewayService_EnableTierFallbackChain(t *testing.T) {
|
|
require.False(t, (*GatewayService)(nil).enableTierFallbackChain())
|
|
require.False(t, (&GatewayService{}).enableTierFallbackChain())
|
|
|
|
cfgOff := &config.Config{}
|
|
cfgOff.Gateway.Scheduling.EnableTierFallbackChain = false
|
|
require.False(t, (&GatewayService{cfg: cfgOff}).enableTierFallbackChain())
|
|
|
|
cfgOn := &config.Config{}
|
|
cfgOn.Gateway.Scheduling.EnableTierFallbackChain = true
|
|
require.True(t, (&GatewayService{cfg: cfgOn}).enableTierFallbackChain())
|
|
}
|
|
|
|
// TestGatewayService_SelectAccountWithTierFallback_PrefersSubscription verifies
|
|
// that when both OAuth (subscription) and APIKey accounts are available, the
|
|
// tier-0 OAuth account is always selected first even if APIKey has higher priority.
|
|
func TestGatewayService_SelectAccountWithTierFallback_PrefersSubscription(t *testing.T) {
|
|
ctx := context.Background()
|
|
|
|
oauthAcc := Account{ID: 91001, Platform: PlatformAnthropic, Type: AccountTypeOAuth, Status: StatusActive, Schedulable: true, Priority: 5}
|
|
apiKeyAcc := Account{ID: 91002, Platform: PlatformAnthropic, Type: AccountTypeAPIKey, Status: StatusActive, Schedulable: true, Priority: 0}
|
|
|
|
repo := &mockAccountRepoForPlatform{
|
|
accounts: []Account{oauthAcc, apiKeyAcc},
|
|
accountsByID: map[int64]*Account{91001: &oauthAcc, 91002: &apiKeyAcc},
|
|
}
|
|
cache := &mockGatewayCacheForPlatform{}
|
|
svc := &GatewayService{accountRepo: repo, cache: cache, cfg: testConfig()}
|
|
|
|
acc, err := svc.selectAccountWithTierFallback(ctx, nil, "", "", nil)
|
|
require.NoError(t, err)
|
|
require.NotNil(t, acc)
|
|
require.Equal(t, int64(91001), acc.ID, "OAuth (tier-0) account should be preferred over APIKey (tier-1)")
|
|
}
|
|
|
|
// TestGatewayService_SelectAccountWithTierFallback_FallsBackToAPIKey verifies
|
|
// that when the subscription tier has no schedulable accounts, the fallback
|
|
// selects an API Key account.
|
|
func TestGatewayService_SelectAccountWithTierFallback_FallsBackToAPIKey(t *testing.T) {
|
|
ctx := context.Background()
|
|
|
|
rateLimitedUntil := time.Now().Add(30 * time.Minute)
|
|
oauthAcc := Account{ID: 92001, Platform: PlatformAnthropic, Type: AccountTypeOAuth, Status: StatusActive, Schedulable: true, RateLimitResetAt: &rateLimitedUntil}
|
|
apiKeyAcc := Account{ID: 92002, Platform: PlatformAnthropic, Type: AccountTypeAPIKey, Status: StatusActive, Schedulable: true}
|
|
|
|
repo := &mockAccountRepoForPlatform{
|
|
accounts: []Account{oauthAcc, apiKeyAcc},
|
|
accountsByID: map[int64]*Account{92001: &oauthAcc, 92002: &apiKeyAcc},
|
|
}
|
|
cache := &mockGatewayCacheForPlatform{}
|
|
svc := &GatewayService{accountRepo: repo, cache: cache, cfg: testConfig()}
|
|
|
|
acc, err := svc.selectAccountWithTierFallback(ctx, nil, "", "", nil)
|
|
require.NoError(t, err)
|
|
require.NotNil(t, acc)
|
|
require.Equal(t, int64(92002), acc.ID, "Should fall back to APIKey when OAuth is rate-limited")
|
|
}
|
|
|
|
// TestGatewayService_SelectAccountWithTierFallback_ExcludesAccounts ensures
|
|
// excluded IDs are respected across all tiers.
|
|
func TestGatewayService_SelectAccountWithTierFallback_ExcludesAccounts(t *testing.T) {
|
|
ctx := context.Background()
|
|
|
|
oauthAcc := Account{ID: 93001, Platform: PlatformAnthropic, Type: AccountTypeOAuth, Status: StatusActive, Schedulable: true}
|
|
apiKeyAcc := Account{ID: 93002, Platform: PlatformAnthropic, Type: AccountTypeAPIKey, Status: StatusActive, Schedulable: true}
|
|
|
|
repo := &mockAccountRepoForPlatform{
|
|
accounts: []Account{oauthAcc, apiKeyAcc},
|
|
accountsByID: map[int64]*Account{93001: &oauthAcc, 93002: &apiKeyAcc},
|
|
}
|
|
cache := &mockGatewayCacheForPlatform{}
|
|
svc := &GatewayService{accountRepo: repo, cache: cache, cfg: testConfig()}
|
|
|
|
excluded := map[int64]struct{}{93001: {}}
|
|
acc, err := svc.selectAccountWithTierFallback(ctx, nil, "", "", excluded)
|
|
require.NoError(t, err)
|
|
require.NotNil(t, acc)
|
|
require.Equal(t, int64(93002), acc.ID, "Excluded OAuth account should cause APIKey fallback")
|
|
}
|
|
|
|
// TestGatewayService_SelectAccountWithTierFallback_NoAccounts verifies that
|
|
// an error is returned when all tiers are empty.
|
|
func TestGatewayService_SelectAccountWithTierFallback_NoAccounts(t *testing.T) {
|
|
ctx := context.Background()
|
|
|
|
repo := &mockAccountRepoForPlatform{accounts: nil, accountsByID: map[int64]*Account{}}
|
|
cache := &mockGatewayCacheForPlatform{}
|
|
svc := &GatewayService{accountRepo: repo, cache: cache, cfg: testConfig()}
|
|
|
|
acc, err := svc.selectAccountWithTierFallback(ctx, nil, "", "", nil)
|
|
require.Error(t, err)
|
|
require.Nil(t, acc)
|
|
}
|
|
|
|
// TestGatewayService_SelectAccountWithTierFallback_BedrockLastResort verifies
|
|
// that Bedrock accounts are only used when subscription and API Key tiers are exhausted.
|
|
func TestGatewayService_SelectAccountWithTierFallback_BedrockLastResort(t *testing.T) {
|
|
ctx := context.Background()
|
|
|
|
rateLimitedUntil := time.Now().Add(30 * time.Minute)
|
|
oauthAcc := Account{ID: 94001, Platform: PlatformAnthropic, Type: AccountTypeOAuth, Status: StatusActive, Schedulable: true, RateLimitResetAt: &rateLimitedUntil}
|
|
apiKeyAcc := Account{ID: 94002, Platform: PlatformAnthropic, Type: AccountTypeAPIKey, Status: StatusActive, Schedulable: true, RateLimitResetAt: &rateLimitedUntil}
|
|
bedrockAcc := Account{ID: 94003, Platform: PlatformAnthropic, Type: AccountTypeBedrock, Status: StatusActive, Schedulable: true}
|
|
|
|
repo := &mockAccountRepoForPlatform{
|
|
accounts: []Account{oauthAcc, apiKeyAcc, bedrockAcc},
|
|
accountsByID: map[int64]*Account{94001: &oauthAcc, 94002: &apiKeyAcc, 94003: &bedrockAcc},
|
|
}
|
|
cache := &mockGatewayCacheForPlatform{}
|
|
svc := &GatewayService{accountRepo: repo, cache: cache, cfg: testConfig()}
|
|
|
|
acc, err := svc.selectAccountWithTierFallback(ctx, nil, "", "", nil)
|
|
require.NoError(t, err)
|
|
require.NotNil(t, acc)
|
|
require.Equal(t, int64(94003), acc.ID, "Bedrock should be selected as last resort")
|
|
}
|