package service import ( "context" "testing" "time" "github.com/Wei-Shaw/sub2api/internal/config" "github.com/stretchr/testify/require" ) func TestOpenAIGatewayService_SelectAccountByPreviousResponseID_Hit(t *testing.T) { ctx := context.Background() groupID := int64(23) account := Account{ ID: 2, Platform: PlatformOpenAI, Type: AccountTypeAPIKey, Status: StatusActive, Schedulable: true, Concurrency: 2, Extra: map[string]any{ "openai_apikey_responses_websockets_v2_enabled": true, }, } cache := &stubGatewayCache{} store := NewOpenAIWSStateStore(cache) cfg := newOpenAIWSV2TestConfig() svc := &OpenAIGatewayService{ accountRepo: stubOpenAIAccountRepo{accounts: []Account{account}}, cache: cache, cfg: cfg, concurrencyService: NewConcurrencyService(stubConcurrencyCache{}), openaiWSStateStore: store, } require.NoError(t, store.BindResponseAccount(ctx, groupID, "resp_prev_1", account.ID, time.Hour)) selection, err := svc.SelectAccountByPreviousResponseID(ctx, &groupID, "resp_prev_1", "gpt-5.1", nil) require.NoError(t, err) require.NotNil(t, selection) require.NotNil(t, selection.Account) require.Equal(t, account.ID, selection.Account.ID) require.True(t, selection.Acquired) if selection.ReleaseFunc != nil { selection.ReleaseFunc() } } func TestOpenAIGatewayService_SelectAccountByPreviousResponseID_Excluded(t *testing.T) { ctx := context.Background() groupID := int64(23) account := Account{ ID: 8, Platform: PlatformOpenAI, Type: AccountTypeAPIKey, Status: StatusActive, Schedulable: true, Concurrency: 1, Extra: map[string]any{ "openai_apikey_responses_websockets_v2_enabled": true, }, } cache := &stubGatewayCache{} store := NewOpenAIWSStateStore(cache) cfg := newOpenAIWSV2TestConfig() svc := &OpenAIGatewayService{ accountRepo: stubOpenAIAccountRepo{accounts: []Account{account}}, cache: cache, cfg: cfg, concurrencyService: NewConcurrencyService(stubConcurrencyCache{}), openaiWSStateStore: store, } require.NoError(t, store.BindResponseAccount(ctx, groupID, "resp_prev_2", account.ID, time.Hour)) selection, err := svc.SelectAccountByPreviousResponseID(ctx, &groupID, "resp_prev_2", "gpt-5.1", map[int64]struct{}{account.ID: {}}) require.NoError(t, err) require.Nil(t, selection) } func TestOpenAIGatewayService_SelectAccountByPreviousResponseID_ForceHTTPIgnored(t *testing.T) { ctx := context.Background() groupID := int64(23) account := Account{ ID: 11, Platform: PlatformOpenAI, Type: AccountTypeAPIKey, Status: StatusActive, Schedulable: true, Concurrency: 1, Extra: map[string]any{ "openai_ws_force_http": true, "responses_websockets_v2_enabled": true, }, } cache := &stubGatewayCache{} store := NewOpenAIWSStateStore(cache) cfg := newOpenAIWSV2TestConfig() svc := &OpenAIGatewayService{ accountRepo: stubOpenAIAccountRepo{accounts: []Account{account}}, cache: cache, cfg: cfg, concurrencyService: NewConcurrencyService(stubConcurrencyCache{}), openaiWSStateStore: store, } require.NoError(t, store.BindResponseAccount(ctx, groupID, "resp_prev_force_http", account.ID, time.Hour)) selection, err := svc.SelectAccountByPreviousResponseID(ctx, &groupID, "resp_prev_force_http", "gpt-5.1", nil) require.NoError(t, err) require.Nil(t, selection, "force_http 场景应忽略 previous_response_id 粘连") } func TestOpenAIGatewayService_SelectAccountByPreviousResponseID_BusyKeepsSticky(t *testing.T) { ctx := context.Background() groupID := int64(23) accounts := []Account{ { ID: 21, Platform: PlatformOpenAI, Type: AccountTypeAPIKey, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 0, Extra: map[string]any{ "openai_apikey_responses_websockets_v2_enabled": true, }, }, { ID: 22, Platform: PlatformOpenAI, Type: AccountTypeAPIKey, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 9, Extra: map[string]any{ "openai_apikey_responses_websockets_v2_enabled": true, }, }, } cache := &stubGatewayCache{} store := NewOpenAIWSStateStore(cache) cfg := newOpenAIWSV2TestConfig() cfg.Gateway.Scheduling.StickySessionMaxWaiting = 2 cfg.Gateway.Scheduling.StickySessionWaitTimeout = 30 * time.Second concurrencyCache := stubConcurrencyCache{ acquireResults: map[int64]bool{ 21: false, // previous_response 命中的账号繁忙 22: true, // 次优账号可用(若回退会命中) }, waitCounts: map[int64]int{ 21: 999, }, } svc := &OpenAIGatewayService{ accountRepo: stubOpenAIAccountRepo{accounts: accounts}, cache: cache, cfg: cfg, concurrencyService: NewConcurrencyService(concurrencyCache), openaiWSStateStore: store, } require.NoError(t, store.BindResponseAccount(ctx, groupID, "resp_prev_busy", 21, time.Hour)) selection, err := svc.SelectAccountByPreviousResponseID(ctx, &groupID, "resp_prev_busy", "gpt-5.1", nil) require.NoError(t, err) require.NotNil(t, selection) require.NotNil(t, selection.Account) require.Equal(t, int64(21), selection.Account.ID, "busy previous_response sticky account should remain selected") require.False(t, selection.Acquired) require.NotNil(t, selection.WaitPlan) require.Equal(t, int64(21), selection.WaitPlan.AccountID) } func newOpenAIWSV2TestConfig() *config.Config { cfg := &config.Config{} cfg.Gateway.OpenAIWS.Enabled = true cfg.Gateway.OpenAIWS.OAuthEnabled = true cfg.Gateway.OpenAIWS.APIKeyEnabled = true cfg.Gateway.OpenAIWS.ResponsesWebsocketsV2 = true cfg.Gateway.OpenAIWS.StickyResponseIDTTLSeconds = 3600 return cfg }