64 lines
1.9 KiB
Go
64 lines
1.9 KiB
Go
//go:build unit
|
|
|
|
package service
|
|
|
|
import (
|
|
"context"
|
|
"testing"
|
|
|
|
"github.com/Wei-Shaw/sub2api/internal/config"
|
|
"github.com/stretchr/testify/require"
|
|
)
|
|
|
|
func TestOpenAIGatewayService_SelectAccountWithScheduler_UsesWSPassthroughSnapshotFlags(t *testing.T) {
|
|
ctx := context.Background()
|
|
groupID := int64(10105)
|
|
account := &Account{
|
|
ID: 35001,
|
|
Platform: PlatformOpenAI,
|
|
Type: AccountTypeOAuth,
|
|
Status: StatusActive,
|
|
Schedulable: true,
|
|
Concurrency: 10,
|
|
Extra: map[string]any{
|
|
"openai_oauth_responses_websockets_v2_mode": OpenAIWSIngressModePassthrough,
|
|
},
|
|
}
|
|
|
|
snapshotCache := &openAISnapshotCacheStub{
|
|
snapshotAccounts: []*Account{account},
|
|
accountsByID: map[int64]*Account{account.ID: account},
|
|
}
|
|
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.ModeRouterV2Enabled = true
|
|
cfg.Gateway.OpenAIWS.IngressModeDefault = OpenAIWSIngressModeCtxPool
|
|
|
|
svc := &OpenAIGatewayService{
|
|
accountRepo: schedulerTestOpenAIAccountRepo{accounts: []Account{*account}},
|
|
cache: &schedulerTestGatewayCache{},
|
|
cfg: cfg,
|
|
rateLimitService: newOpenAIAdvancedSchedulerRateLimitService("true"),
|
|
schedulerSnapshot: &SchedulerSnapshotService{cache: snapshotCache},
|
|
concurrencyService: NewConcurrencyService(schedulerTestConcurrencyCache{}),
|
|
}
|
|
|
|
selection, decision, err := svc.SelectAccountWithScheduler(
|
|
ctx,
|
|
&groupID,
|
|
"",
|
|
"session_hash_ws_passthrough",
|
|
"gpt-5.1",
|
|
nil,
|
|
OpenAIUpstreamTransportResponsesWebsocketV2,
|
|
)
|
|
require.NoError(t, err)
|
|
require.NotNil(t, selection)
|
|
require.NotNil(t, selection.Account)
|
|
require.Equal(t, account.ID, selection.Account.ID)
|
|
require.Equal(t, openAIAccountScheduleLayerLoadBalance, decision.Layer)
|
|
}
|