| 123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502 |
- package service
- import (
- "context"
- "errors"
- "testing"
- "github.com/Wei-Shaw/sub2api/internal/config"
- "github.com/stretchr/testify/require"
- )
- // TestOpenAIGatewayService_SelectAccountWithScheduler_CompactPrefersSupportedOverUnknown
- // 验证 compact 调度时显式支持 (tier=2) 优先于未探测 (tier=1)。
- func TestOpenAIGatewayService_SelectAccountWithScheduler_CompactPrefersSupportedOverUnknown(t *testing.T) {
- resetOpenAIAdvancedSchedulerSettingCacheForTest()
- ctx := context.Background()
- groupID := int64(91001)
- accounts := []Account{
- {
- ID: 71001,
- Platform: PlatformOpenAI,
- Type: AccountTypeAPIKey,
- Status: StatusActive,
- Schedulable: true,
- Concurrency: 1,
- Priority: 0,
- Extra: map[string]any{}, // unknown
- },
- {
- ID: 71002,
- Platform: PlatformOpenAI,
- Type: AccountTypeAPIKey,
- Status: StatusActive,
- Schedulable: true,
- Concurrency: 1,
- Priority: 0,
- Extra: map[string]any{"openai_compact_supported": true}, // tier=2
- },
- }
- cfg := &config.Config{}
- cfg.Gateway.Scheduling.LoadBatchEnabled = false
- svc := &OpenAIGatewayService{
- accountRepo: schedulerTestOpenAIAccountRepo{accounts: accounts},
- cache: &schedulerTestGatewayCache{},
- cfg: cfg,
- concurrencyService: NewConcurrencyService(schedulerTestConcurrencyCache{}),
- }
- selection, _, err := svc.SelectAccountWithScheduler(
- ctx,
- &groupID,
- "",
- "",
- "gpt-5.4",
- nil,
- OpenAIUpstreamTransportAny,
- true,
- )
- require.NoError(t, err)
- require.NotNil(t, selection)
- require.NotNil(t, selection.Account)
- require.Equal(t, int64(71002), selection.Account.ID, "compact-supported account should win over unknown")
- }
- // TestOpenAIGatewayService_SelectAccountWithScheduler_CompactRejectsExplicitlyUnsupported
- // 验证 force_off / 已探测不支持 (tier=0) 的账号不会被 compact 请求选中。
- func TestOpenAIGatewayService_SelectAccountWithScheduler_CompactRejectsExplicitlyUnsupported(t *testing.T) {
- resetOpenAIAdvancedSchedulerSettingCacheForTest()
- ctx := context.Background()
- groupID := int64(91002)
- accounts := []Account{
- {
- ID: 71010,
- Platform: PlatformOpenAI,
- Type: AccountTypeAPIKey,
- Status: StatusActive,
- Schedulable: true,
- Concurrency: 1,
- Priority: 0,
- Extra: map[string]any{"openai_compact_mode": OpenAICompactModeForceOff},
- },
- {
- ID: 71011,
- Platform: PlatformOpenAI,
- Type: AccountTypeAPIKey,
- Status: StatusActive,
- Schedulable: true,
- Concurrency: 1,
- Priority: 0,
- Extra: map[string]any{"openai_compact_supported": false},
- },
- }
- cfg := &config.Config{}
- cfg.Gateway.Scheduling.LoadBatchEnabled = false
- svc := &OpenAIGatewayService{
- accountRepo: schedulerTestOpenAIAccountRepo{accounts: accounts},
- cache: &schedulerTestGatewayCache{},
- cfg: cfg,
- concurrencyService: NewConcurrencyService(schedulerTestConcurrencyCache{}),
- }
- selection, _, err := svc.SelectAccountWithScheduler(
- ctx,
- &groupID,
- "",
- "",
- "gpt-5.4",
- nil,
- OpenAIUpstreamTransportAny,
- true,
- )
- require.Error(t, err)
- require.True(t, errors.Is(err, ErrNoAvailableCompactAccounts), "compact-only accounts should rejected explicitly unsupported and return compact error")
- require.Nil(t, selection)
- }
- func newOpenAICompactionSchedulerTestService(accounts []Account, advanced bool) *OpenAIGatewayService {
- cfg := &config.Config{}
- cfg.Gateway.Scheduling.LoadBatchEnabled = false
- svc := &OpenAIGatewayService{
- accountRepo: schedulerTestOpenAIAccountRepo{accounts: accounts},
- cache: &schedulerTestGatewayCache{},
- cfg: cfg,
- concurrencyService: NewConcurrencyService(schedulerTestConcurrencyCache{}),
- }
- if advanced {
- svc.rateLimitService = newOpenAIAdvancedSchedulerRateLimitService("true")
- }
- return svc
- }
- func selectOpenAICompactionSchedulerTestAccount(t *testing.T, svc *OpenAIGatewayService, groupID int64, requireCompact bool) (*AccountSelectionResult, error) {
- t.Helper()
- selection, _, err := svc.SelectAccountWithSchedulerForCapability(
- context.Background(),
- &groupID,
- "",
- "",
- "gpt-5.6-sol",
- nil,
- OpenAIUpstreamTransportAny,
- OpenAIEndpointCapabilityResponses,
- requireCompact,
- false,
- false,
- )
- return selection, err
- }
- func TestOpenAIGatewayService_SelectAccountWithScheduler_NativeCompactionIgnoresLegacyCompactProbe(t *testing.T) {
- for _, advanced := range []bool{false, true} {
- t.Run(map[bool]string{false: "legacy_scheduler", true: "advanced_scheduler"}[advanced], func(t *testing.T) {
- resetOpenAIAdvancedSchedulerSettingCacheForTest()
- svc := newOpenAICompactionSchedulerTestService([]Account{{
- ID: 71012,
- Platform: PlatformOpenAI,
- Type: AccountTypeAPIKey,
- Status: StatusActive,
- Schedulable: true,
- Concurrency: 1,
- Extra: map[string]any{
- "openai_compact_supported": false,
- "openai_responses_supported": true,
- },
- }}, advanced)
- selection, err := selectOpenAICompactionSchedulerTestAccount(t, svc, 91007, false)
- require.NoError(t, err)
- require.NotNil(t, selection)
- require.Equal(t, int64(71012), selection.Account.ID)
- })
- }
- }
- func TestOpenAIGatewayService_SelectAccountWithScheduler_NativeCompactionAllowsForceOff(t *testing.T) {
- for _, advanced := range []bool{false, true} {
- t.Run(map[bool]string{false: "legacy_scheduler", true: "advanced_scheduler"}[advanced], func(t *testing.T) {
- resetOpenAIAdvancedSchedulerSettingCacheForTest()
- svc := newOpenAICompactionSchedulerTestService([]Account{{
- ID: 71013,
- Platform: PlatformOpenAI,
- Type: AccountTypeAPIKey,
- Status: StatusActive,
- Schedulable: true,
- Concurrency: 1,
- Extra: map[string]any{
- "openai_compact_mode": OpenAICompactModeForceOff,
- "openai_responses_supported": true,
- },
- }}, advanced)
- selection, err := selectOpenAICompactionSchedulerTestAccount(t, svc, 91008, false)
- require.NoError(t, err)
- require.NotNil(t, selection)
- require.Equal(t, int64(71013), selection.Account.ID)
- })
- }
- }
- func TestOpenAIGatewayService_SelectAccountWithScheduler_NativeCompactionRequiresResponsesCapability(t *testing.T) {
- for _, advanced := range []bool{false, true} {
- t.Run(map[bool]string{false: "legacy_scheduler", true: "advanced_scheduler"}[advanced], func(t *testing.T) {
- resetOpenAIAdvancedSchedulerSettingCacheForTest()
- svc := newOpenAICompactionSchedulerTestService([]Account{{
- ID: 71014,
- Platform: PlatformOpenAI,
- Type: AccountTypeAPIKey,
- Status: StatusActive,
- Schedulable: true,
- Concurrency: 1,
- Extra: map[string]any{
- "openai_compact_mode": OpenAICompactModeForceOn,
- "openai_responses_supported": false,
- },
- }}, advanced)
- selection, err := selectOpenAICompactionSchedulerTestAccount(t, svc, 91009, false)
- require.ErrorIs(t, err, ErrNoAvailableAccounts)
- require.NotErrorIs(t, err, ErrNoAvailableCompactAccounts)
- require.NotContains(t, err.Error(), "/responses/compact")
- require.Nil(t, selection)
- })
- }
- }
- func TestOpenAIGatewayService_SelectAccountWithScheduler_LegacyCompactionKeepsCompactEligibility(t *testing.T) {
- for _, advanced := range []bool{false, true} {
- t.Run(map[bool]string{false: "legacy_scheduler", true: "advanced_scheduler"}[advanced], func(t *testing.T) {
- resetOpenAIAdvancedSchedulerSettingCacheForTest()
- svc := newOpenAICompactionSchedulerTestService([]Account{
- {
- ID: 71015,
- Platform: PlatformOpenAI,
- Type: AccountTypeAPIKey,
- Status: StatusActive,
- Schedulable: true,
- Concurrency: 1,
- Extra: map[string]any{
- "openai_compact_supported": false,
- "openai_responses_supported": true,
- },
- },
- {
- ID: 71016,
- Platform: PlatformOpenAI,
- Type: AccountTypeAPIKey,
- Status: StatusActive,
- Schedulable: true,
- Concurrency: 1,
- Extra: map[string]any{
- "openai_compact_mode": OpenAICompactModeForceOff,
- "openai_responses_supported": true,
- },
- },
- }, advanced)
- selection, err := selectOpenAICompactionSchedulerTestAccount(t, svc, 91010, true)
- require.ErrorIs(t, err, ErrNoAvailableCompactAccounts)
- require.Contains(t, err.Error(), "/responses/compact")
- require.Nil(t, selection)
- })
- }
- }
- // TestOpenAIGatewayService_SelectAccountWithScheduler_CompactRequiresResponsesCapability
- // prevents a confirmed Chat Completions-only API key from silently dropping the
- // compaction trigger through the Responses-to-Chat fallback.
- func TestOpenAIGatewayService_SelectAccountWithScheduler_CompactRequiresResponsesCapability(t *testing.T) {
- resetOpenAIAdvancedSchedulerSettingCacheForTest()
- ctx := context.Background()
- groupID := int64(91005)
- accounts := []Account{{
- ID: 71050,
- Platform: PlatformOpenAI,
- Type: AccountTypeAPIKey,
- Status: StatusActive,
- Schedulable: true,
- Concurrency: 1,
- Extra: map[string]any{
- "openai_compact_supported": true,
- "openai_responses_supported": false,
- },
- }}
- cfg := &config.Config{}
- cfg.Gateway.Scheduling.LoadBatchEnabled = false
- svc := &OpenAIGatewayService{
- accountRepo: schedulerTestOpenAIAccountRepo{accounts: accounts},
- cache: &schedulerTestGatewayCache{},
- cfg: cfg,
- concurrencyService: NewConcurrencyService(schedulerTestConcurrencyCache{}),
- }
- selection, _, err := svc.SelectAccountWithSchedulerForCapability(
- ctx,
- &groupID,
- "",
- "",
- "gpt-5.4",
- nil,
- OpenAIUpstreamTransportAny,
- OpenAIEndpointCapabilityResponses,
- true,
- false,
- false,
- )
- require.Error(t, err)
- require.True(t, errors.Is(err, ErrNoAvailableAccounts))
- require.Nil(t, selection)
- }
- func TestOpenAIGatewayService_SelectAccountWithScheduler_CompactSkipsChatOnlyAccount(t *testing.T) {
- resetOpenAIAdvancedSchedulerSettingCacheForTest()
- ctx := context.Background()
- groupID := int64(91006)
- accounts := []Account{
- {
- ID: 71060,
- Platform: PlatformOpenAI,
- Type: AccountTypeAPIKey,
- Status: StatusActive,
- Schedulable: true,
- Concurrency: 1,
- Priority: 10,
- Extra: map[string]any{
- "openai_compact_supported": true,
- "openai_responses_supported": false,
- },
- },
- {
- ID: 71061,
- Platform: PlatformOpenAI,
- Type: AccountTypeAPIKey,
- Status: StatusActive,
- Schedulable: true,
- Concurrency: 1,
- Priority: 0,
- Extra: map[string]any{
- "openai_compact_supported": true,
- "openai_responses_supported": true,
- },
- },
- }
- cfg := &config.Config{}
- cfg.Gateway.Scheduling.LoadBatchEnabled = false
- svc := &OpenAIGatewayService{
- accountRepo: schedulerTestOpenAIAccountRepo{accounts: accounts},
- cache: &schedulerTestGatewayCache{},
- cfg: cfg,
- concurrencyService: NewConcurrencyService(schedulerTestConcurrencyCache{}),
- }
- selection, _, err := svc.SelectAccountWithSchedulerForCapability(
- ctx,
- &groupID,
- "",
- "",
- "gpt-5.6-sol",
- nil,
- OpenAIUpstreamTransportAny,
- OpenAIEndpointCapabilityResponses,
- true,
- false,
- false,
- )
- require.NoError(t, err)
- require.NotNil(t, selection)
- require.NotNil(t, selection.Account)
- require.Equal(t, int64(71061), selection.Account.ID)
- }
- // TestOpenAIGatewayService_SelectAccountWithScheduler_CompactFallsBackToUnknown
- // 验证当没有"已知支持"账号时,compact 请求会回退到"未探测"账号。
- func TestOpenAIGatewayService_SelectAccountWithScheduler_CompactFallsBackToUnknown(t *testing.T) {
- resetOpenAIAdvancedSchedulerSettingCacheForTest()
- ctx := context.Background()
- groupID := int64(91003)
- accounts := []Account{
- {
- ID: 71020,
- Platform: PlatformOpenAI,
- Type: AccountTypeAPIKey,
- Status: StatusActive,
- Schedulable: true,
- Concurrency: 1,
- Priority: 0,
- Extra: map[string]any{"openai_compact_supported": false}, // tier=0
- },
- {
- ID: 71021,
- Platform: PlatformOpenAI,
- Type: AccountTypeAPIKey,
- Status: StatusActive,
- Schedulable: true,
- Concurrency: 1,
- Priority: 0,
- Extra: map[string]any{}, // unknown -> tier=1
- },
- }
- cfg := &config.Config{}
- cfg.Gateway.Scheduling.LoadBatchEnabled = false
- svc := &OpenAIGatewayService{
- accountRepo: schedulerTestOpenAIAccountRepo{accounts: accounts},
- cache: &schedulerTestGatewayCache{},
- cfg: cfg,
- concurrencyService: NewConcurrencyService(schedulerTestConcurrencyCache{}),
- }
- selection, _, err := svc.SelectAccountWithScheduler(
- ctx,
- &groupID,
- "",
- "",
- "gpt-5.4",
- nil,
- OpenAIUpstreamTransportAny,
- true,
- )
- require.NoError(t, err)
- require.NotNil(t, selection)
- require.NotNil(t, selection.Account)
- require.Equal(t, int64(71021), selection.Account.ID, "unknown account should be picked when no supported account available")
- }
- // TestOpenAIGatewayService_SelectAccountWithScheduler_CompactAllowsGrok verifies
- // that OpenAI-compatible compact routing does not reject Grok accounts.
- func TestOpenAIGatewayService_SelectAccountWithScheduler_CompactAllowsGrok(t *testing.T) {
- resetOpenAIAdvancedSchedulerSettingCacheForTest()
- ctx := context.Background()
- groupID := int64(91004)
- accounts := []Account{
- {
- ID: 71030,
- Platform: PlatformGrok,
- Type: AccountTypeOAuth,
- Status: StatusActive,
- Schedulable: true,
- Concurrency: 1,
- Priority: 0,
- Credentials: map[string]any{
- "model_mapping": map[string]any{"grok-4.5": "grok-4.5"},
- },
- },
- }
- cfg := &config.Config{}
- cfg.Gateway.Scheduling.LoadBatchEnabled = false
- svc := &OpenAIGatewayService{
- accountRepo: schedulerTestOpenAIAccountRepo{accounts: accounts},
- cache: &schedulerTestGatewayCache{},
- cfg: cfg,
- concurrencyService: NewConcurrencyService(schedulerTestConcurrencyCache{}),
- }
- selection, _, err := svc.SelectAccountWithSchedulerForCapability(
- ctx,
- &groupID,
- "",
- "",
- "grok-4.5",
- nil,
- OpenAIUpstreamTransportAny,
- OpenAIEndpointCapabilityChatCompletions,
- true,
- false,
- true,
- PlatformGrok,
- )
- require.NoError(t, err)
- require.NotNil(t, selection)
- require.NotNil(t, selection.Account)
- require.Equal(t, int64(71030), selection.Account.ID)
- }
- // TestOpenAICompactSupportTier 验证 tier 分类逻辑。
- func TestOpenAICompactSupportTier(t *testing.T) {
- tests := []struct {
- name string
- account *Account
- want int
- }{
- {name: "nil", account: nil, want: 0},
- {name: "non openai", account: &Account{Platform: PlatformAnthropic}, want: 0},
- {name: "grok", account: &Account{Platform: PlatformGrok}, want: 2},
- {name: "openai unknown", account: &Account{Platform: PlatformOpenAI, Extra: map[string]any{}}, want: 1},
- {name: "openai supported", account: &Account{Platform: PlatformOpenAI, Extra: map[string]any{"openai_compact_supported": true}}, want: 2},
- {name: "openai unsupported", account: &Account{Platform: PlatformOpenAI, Extra: map[string]any{"openai_compact_supported": false}}, want: 0},
- {name: "force on", account: &Account{Platform: PlatformOpenAI, Extra: map[string]any{"openai_compact_mode": OpenAICompactModeForceOn}}, want: 2},
- {name: "force off overrides probe true", account: &Account{Platform: PlatformOpenAI, Extra: map[string]any{"openai_compact_mode": OpenAICompactModeForceOff, "openai_compact_supported": true}}, want: 0},
- }
- for _, tt := range tests {
- t.Run(tt.name, func(t *testing.T) {
- if got := openAICompactSupportTier(tt.account); got != tt.want {
- t.Fatalf("openAICompactSupportTier(...) = %d, want %d", got, tt.want)
- }
- })
- }
- }
|