| 123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520 |
- package service
- import (
- "context"
- "encoding/json"
- "errors"
- "testing"
- "github.com/Wei-Shaw/sub2api/internal/config"
- "github.com/Wei-Shaw/sub2api/internal/pkg/ctxkey"
- "github.com/stretchr/testify/require"
- "github.com/tidwall/gjson"
- )
- type openAIFastPolicyRepoStub struct {
- values map[string]string
- }
- func (s *openAIFastPolicyRepoStub) Get(ctx context.Context, key string) (*Setting, error) {
- panic("unexpected Get call")
- }
- func (s *openAIFastPolicyRepoStub) GetValue(ctx context.Context, key string) (string, error) {
- if v, ok := s.values[key]; ok {
- return v, nil
- }
- return "", ErrSettingNotFound
- }
- func (s *openAIFastPolicyRepoStub) Set(ctx context.Context, key, value string) error {
- if s.values == nil {
- s.values = map[string]string{}
- }
- s.values[key] = value
- return nil
- }
- func (s *openAIFastPolicyRepoStub) GetMultiple(ctx context.Context, keys []string) (map[string]string, error) {
- result := make(map[string]string, len(keys))
- for _, key := range keys {
- if value, ok := s.values[key]; ok {
- result[key] = value
- }
- }
- return result, nil
- }
- func (s *openAIFastPolicyRepoStub) SetMultiple(ctx context.Context, settings map[string]string) error {
- panic("unexpected SetMultiple call")
- }
- func (s *openAIFastPolicyRepoStub) GetAll(ctx context.Context) (map[string]string, error) {
- panic("unexpected GetAll call")
- }
- func (s *openAIFastPolicyRepoStub) Delete(ctx context.Context, key string) error {
- panic("unexpected Delete call")
- }
- func newOpenAIGatewayServiceWithSettings(t *testing.T, settings *OpenAIFastPolicySettings) *OpenAIGatewayService {
- t.Helper()
- repo := &openAIFastPolicyRepoStub{values: map[string]string{}}
- if settings != nil {
- raw, err := json.Marshal(settings)
- require.NoError(t, err)
- repo.values[SettingKeyOpenAIFastPolicySettings] = string(raw)
- }
- return &OpenAIGatewayService{
- settingService: NewSettingService(repo, &config.Config{}),
- }
- }
- func openAIFastFilterPriorityPolicy() *OpenAIFastPolicySettings {
- return &OpenAIFastPolicySettings{
- Rules: []OpenAIFastPolicyRule{{
- ServiceTier: OpenAIFastTierPriority,
- Action: BetaPolicyActionFilter,
- Scope: BetaPolicyScopeAll,
- ModelWhitelist: []string{},
- FallbackAction: BetaPolicyActionPass,
- }},
- }
- }
- func TestEvaluateOpenAIFastPolicy_DefaultPassesKnownTiers(t *testing.T) {
- require.Empty(t, DefaultOpenAIFastPolicySettings().Rules, "default policy must not rewrite service_tier unless admin configured rules")
- svc := newOpenAIGatewayServiceWithSettings(t, DefaultOpenAIFastPolicySettings())
- account := &Account{Platform: PlatformOpenAI, Type: AccountTypeAPIKey}
- action, _ := svc.evaluateOpenAIFastPolicy(context.Background(), account, "gpt-5.5", OpenAIFastTierPriority)
- require.Equal(t, BetaPolicyActionPass, action)
- action, _ = svc.evaluateOpenAIFastPolicy(context.Background(), account, "gpt-5.5-turbo", OpenAIFastTierPriority)
- require.Equal(t, BetaPolicyActionPass, action)
- action, _ = svc.evaluateOpenAIFastPolicy(context.Background(), account, "gpt-4", OpenAIFastTierPriority)
- require.Equal(t, BetaPolicyActionPass, action)
- action, _ = svc.evaluateOpenAIFastPolicy(context.Background(), account, "gpt-5.5", OpenAIFastTierFlex)
- require.Equal(t, BetaPolicyActionPass, action)
- action, _ = svc.evaluateOpenAIFastPolicy(context.Background(), account, "gpt-5.6-sol", OpenAIFastTierUltrafast)
- require.Equal(t, BetaPolicyActionPass, action)
- // empty tier → pass
- action, _ = svc.evaluateOpenAIFastPolicy(context.Background(), account, "gpt-5.5", "")
- require.Equal(t, BetaPolicyActionPass, action)
- }
- func TestEvaluateOpenAIFastPolicy_BlockRuleCarriesMessage(t *testing.T) {
- settings := &OpenAIFastPolicySettings{
- Rules: []OpenAIFastPolicyRule{{
- ServiceTier: OpenAIFastTierPriority,
- Action: BetaPolicyActionBlock,
- Scope: BetaPolicyScopeAll,
- ErrorMessage: "fast mode is not allowed",
- ModelWhitelist: []string{"gpt-5.5"},
- FallbackAction: BetaPolicyActionPass,
- }},
- }
- svc := newOpenAIGatewayServiceWithSettings(t, settings)
- account := &Account{Platform: PlatformOpenAI, Type: AccountTypeAPIKey}
- action, msg := svc.evaluateOpenAIFastPolicy(context.Background(), account, "gpt-5.5", OpenAIFastTierPriority)
- require.Equal(t, BetaPolicyActionBlock, action)
- require.Equal(t, "fast mode is not allowed", msg)
- }
- func TestEvaluateOpenAIFastPolicy_ScopeFiltersOAuth(t *testing.T) {
- settings := &OpenAIFastPolicySettings{
- Rules: []OpenAIFastPolicyRule{{
- ServiceTier: OpenAIFastTierAny,
- Action: BetaPolicyActionFilter,
- Scope: BetaPolicyScopeOAuth,
- }},
- }
- svc := newOpenAIGatewayServiceWithSettings(t, settings)
- // OAuth account → rule matches
- oauthAccount := &Account{Platform: PlatformOpenAI, Type: AccountTypeOAuth}
- action, _ := svc.evaluateOpenAIFastPolicy(context.Background(), oauthAccount, "gpt-4", OpenAIFastTierPriority)
- require.Equal(t, BetaPolicyActionFilter, action)
- // API Key account → rule skipped → pass
- apiKeyAccount := &Account{Platform: PlatformOpenAI, Type: AccountTypeAPIKey}
- action, _ = svc.evaluateOpenAIFastPolicy(context.Background(), apiKeyAccount, "gpt-4", OpenAIFastTierPriority)
- require.Equal(t, BetaPolicyActionPass, action)
- }
- func TestEvaluateOpenAIFastPolicy_UserScopedRuleOverridesGlobalRule(t *testing.T) {
- settings := &OpenAIFastPolicySettings{
- Rules: []OpenAIFastPolicyRule{
- {
- ServiceTier: OpenAIFastTierPriority,
- Action: BetaPolicyActionFilter,
- Scope: BetaPolicyScopeAll,
- },
- {
- ServiceTier: OpenAIFastTierPriority,
- Action: BetaPolicyActionPass,
- Scope: BetaPolicyScopeAll,
- UserIDs: []int64{42},
- },
- },
- }
- svc := newOpenAIGatewayServiceWithSettings(t, settings)
- account := &Account{Platform: PlatformOpenAI, Type: AccountTypeAPIKey}
- allowedUserCtx := context.WithValue(context.Background(), ctxkey.UserID, int64(42))
- action, _ := svc.evaluateOpenAIFastPolicy(allowedUserCtx, account, "gpt-5.5", OpenAIFastTierPriority)
- require.Equal(t, BetaPolicyActionPass, action)
- otherUserCtx := context.WithValue(context.Background(), ctxkey.UserID, int64(43))
- action, _ = svc.evaluateOpenAIFastPolicy(otherUserCtx, account, "gpt-5.5", OpenAIFastTierPriority)
- require.Equal(t, BetaPolicyActionFilter, action)
- action, _ = svc.evaluateOpenAIFastPolicy(context.Background(), account, "gpt-5.5", OpenAIFastTierPriority)
- require.Equal(t, BetaPolicyActionFilter, action)
- }
- func TestApplyOpenAIFastPolicyToBody_DefaultPassesPriorityAndFast(t *testing.T) {
- svc := newOpenAIGatewayServiceWithSettings(t, DefaultOpenAIFastPolicySettings())
- account := &Account{Platform: PlatformOpenAI, Type: AccountTypeAPIKey}
- body := []byte(`{"model":"gpt-5.5","service_tier":"priority","messages":[]}`)
- updated, err := svc.applyOpenAIFastPolicyToBody(context.Background(), account, "gpt-5.5", body)
- require.NoError(t, err)
- require.Equal(t, string(body), string(updated))
- body = []byte(`{"model":"gpt-5.5","service_tier":"fast"}`)
- updated, err = svc.applyOpenAIFastPolicyToBody(context.Background(), account, "gpt-5.5", body)
- require.NoError(t, err)
- require.Equal(t, "priority", gjson.GetBytes(updated, "service_tier").String())
- body = []byte(`{"model":"gpt-4","service_tier":"priority"}`)
- updated, err = svc.applyOpenAIFastPolicyToBody(context.Background(), account, "gpt-4", body)
- require.NoError(t, err)
- require.Equal(t, string(body), string(updated))
- // No service_tier → no-op
- body = []byte(`{"model":"gpt-5.5"}`)
- updated, err = svc.applyOpenAIFastPolicyToBody(context.Background(), account, "gpt-5.5", body)
- require.NoError(t, err)
- require.Equal(t, string(body), string(updated))
- }
- func TestApplyOpenAIFastPolicyToBody_ExplicitFilterRemovesField(t *testing.T) {
- svc := newOpenAIGatewayServiceWithSettings(t, openAIFastFilterPriorityPolicy())
- account := &Account{Platform: PlatformOpenAI, Type: AccountTypeAPIKey}
- body := []byte(`{"model":"gpt-5.5","service_tier":"priority","messages":[]}`)
- updated, err := svc.applyOpenAIFastPolicyToBody(context.Background(), account, "gpt-5.5", body)
- require.NoError(t, err)
- require.NotContains(t, string(updated), `"service_tier"`)
- body = []byte(`{"model":"gpt-5.5","service_tier":"fast"}`)
- updated, err = svc.applyOpenAIFastPolicyToBody(context.Background(), account, "gpt-5.5", body)
- require.NoError(t, err)
- require.NotContains(t, string(updated), `"service_tier"`)
- }
- func TestApplyOpenAIFastPolicyToBody_UserScopedRuleOverridesGlobalRule(t *testing.T) {
- settings := &OpenAIFastPolicySettings{
- Rules: []OpenAIFastPolicyRule{
- {
- ServiceTier: OpenAIFastTierPriority,
- Action: BetaPolicyActionFilter,
- Scope: BetaPolicyScopeAll,
- },
- {
- ServiceTier: OpenAIFastTierPriority,
- Action: BetaPolicyActionPass,
- Scope: BetaPolicyScopeAll,
- UserIDs: []int64{42},
- },
- },
- }
- svc := newOpenAIGatewayServiceWithSettings(t, settings)
- account := &Account{Platform: PlatformOpenAI, Type: AccountTypeAPIKey}
- body := []byte(`{"model":"gpt-5.5","service_tier":"priority"}`)
- allowedUserCtx := context.WithValue(context.Background(), ctxkey.UserID, int64(42))
- updated, err := svc.applyOpenAIFastPolicyToBody(allowedUserCtx, account, "gpt-5.5", body)
- require.NoError(t, err)
- require.Equal(t, "priority", gjson.GetBytes(updated, "service_tier").String())
- otherUserCtx := context.WithValue(context.Background(), ctxkey.UserID, int64(43))
- updated, err = svc.applyOpenAIFastPolicyToBody(otherUserCtx, account, "gpt-5.5", body)
- require.NoError(t, err)
- require.NotContains(t, string(updated), `"service_tier"`)
- }
- func TestApplyOpenAIFastPolicyToBody_PriorityFilterLeavesUltrafast(t *testing.T) {
- svc := newOpenAIGatewayServiceWithSettings(t, openAIFastFilterPriorityPolicy())
- account := &Account{Platform: PlatformOpenAI, Type: AccountTypeAPIKey}
- body := []byte(`{"model":"gpt-5.6-sol","service_tier":"ultrafast"}`)
- updated, err := svc.applyOpenAIFastPolicyToBody(context.Background(), account, "gpt-5.6-sol", body)
- require.NoError(t, err)
- require.Equal(t, OpenAIFastTierUltrafast, gjson.GetBytes(updated, "service_tier").String())
- }
- func TestApplyOpenAIFastPolicyToBody_ForcePriorityRewritesKnownTier(t *testing.T) {
- settings := &OpenAIFastPolicySettings{
- Rules: []OpenAIFastPolicyRule{{
- ServiceTier: OpenAIFastTierAny,
- Action: OpenAIFastPolicyActionForcePriority,
- Scope: BetaPolicyScopeAll,
- }},
- }
- svc := newOpenAIGatewayServiceWithSettings(t, settings)
- account := &Account{Platform: PlatformOpenAI, Type: AccountTypeAPIKey}
- for _, tier := range []string{"flex", "auto", "default", "scale", "fast", "priority", "ultrafast"} {
- body := []byte(`{"model":"gpt-5.5","service_tier":"` + tier + `"}`)
- updated, err := svc.applyOpenAIFastPolicyToBody(context.Background(), account, "gpt-5.5", body)
- require.NoError(t, err)
- require.Equal(t, OpenAIFastTierPriority, gjson.GetBytes(updated, "service_tier").String(),
- "tier %q should be forced to priority", tier)
- }
- }
- func TestApplyOpenAIFastPolicyToBody_GroupForceInjectsAndOverridesTier(t *testing.T) {
- svc := newOpenAIGatewayServiceWithSettings(t, DefaultOpenAIFastPolicySettings())
- account := &Account{Platform: PlatformOpenAI, Type: AccountTypeOAuth}
- ctx := context.WithValue(context.Background(), ctxkey.Group, &Group{
- ID: 7, Platform: PlatformOpenAI, Status: StatusActive, Hydrated: true, ForceOpenAIFast: true,
- })
- for _, body := range [][]byte{
- []byte(`{"model":"gpt-5.6-sol","input":"hi"}`),
- []byte(`{"model":"gpt-5.6-sol","service_tier":"default"}`),
- []byte(`{"model":"gpt-5.6-sol","service_tier":"flex"}`),
- []byte(`{"model":"gpt-5.6-sol","service_tier":"client-unknown"}`),
- } {
- updated, err := svc.applyOpenAIFastPolicyToBody(ctx, account, "gpt-5.6-sol", body)
- require.NoError(t, err)
- require.Equal(t, OpenAIFastTierPriority, gjson.GetBytes(updated, "service_tier").String())
- }
- }
- func TestApplyOpenAIFastPolicyToBody_GroupForceStillHonorsGlobalPolicy(t *testing.T) {
- svc := newOpenAIGatewayServiceWithSettings(t, openAIFastFilterPriorityPolicy())
- account := &Account{Platform: PlatformOpenAI, Type: AccountTypeOAuth}
- ctx := context.WithValue(context.Background(), ctxkey.Group, &Group{
- ID: 7, Platform: PlatformOpenAI, Status: StatusActive, Hydrated: true, ForceOpenAIFast: true,
- })
- updated, err := svc.applyOpenAIFastPolicyToBody(ctx, account, "gpt-5.6-sol", []byte(`{"model":"gpt-5.6-sol"}`))
- require.NoError(t, err)
- require.False(t, gjson.GetBytes(updated, "service_tier").Exists(),
- "the global filter remains authoritative after the group requests priority")
- }
- func TestApplyOpenAIFastPolicyToBody_GroupForceRequiresHydratedGroup(t *testing.T) {
- svc := newOpenAIGatewayServiceWithSettings(t, DefaultOpenAIFastPolicySettings())
- account := &Account{Platform: PlatformOpenAI, Type: AccountTypeOAuth}
- body := []byte(`{"model":"gpt-5.6-sol"}`)
- ctx := context.WithValue(context.Background(), ctxkey.Group, &Group{
- ID: 7, Platform: PlatformOpenAI, Status: StatusActive, ForceOpenAIFast: true,
- })
- updated, err := svc.applyOpenAIFastPolicyToBody(ctx, account, "gpt-5.6-sol", body)
- require.NoError(t, err)
- require.Equal(t, string(body), string(updated))
- }
- func TestApplyOpenAIFastPolicyToBody_GroupForceOnlyTargetsOpenAIAccounts(t *testing.T) {
- svc := newOpenAIGatewayServiceWithSettings(t, DefaultOpenAIFastPolicySettings())
- body := []byte(`{"model":"grok-4.1"}`)
- ctx := context.WithValue(context.Background(), ctxkey.Group, &Group{
- ID: 7, Platform: PlatformComposite, Status: StatusActive, Hydrated: true, ForceOpenAIFast: true,
- })
- updated, err := svc.applyOpenAIFastPolicyToBody(
- ctx,
- &Account{Platform: PlatformGrok, Type: AccountTypeOAuth},
- "grok-4.1",
- body,
- )
- require.NoError(t, err)
- require.Equal(t, string(body), string(updated))
- }
- func TestApplyOpenAIFastPolicyToBody_GroupForceRequiresSupportedGroupPlatform(t *testing.T) {
- svc := newOpenAIGatewayServiceWithSettings(t, DefaultOpenAIFastPolicySettings())
- body := []byte(`{"model":"gpt-5.6-sol"}`)
- ctx := context.WithValue(context.Background(), ctxkey.Group, &Group{
- ID: 7, Platform: PlatformAnthropic, Status: StatusActive, Hydrated: true, ForceOpenAIFast: true,
- })
- updated, err := svc.applyOpenAIFastPolicyToBody(
- ctx,
- &Account{Platform: PlatformOpenAI, Type: AccountTypeOAuth},
- "gpt-5.6-sol",
- body,
- )
- require.NoError(t, err)
- require.Equal(t, string(body), string(updated))
- }
- // TestApplyOpenAIFastPolicyToBody_OfficialTiersBypassDefaultRule 验证默认配置
- // 下客户端显式发送的 OpenAI 官方合法 tier 能透传到上游而不被静默剥离。
- func TestApplyOpenAIFastPolicyToBody_OfficialTiersBypassDefaultRule(t *testing.T) {
- svc := newOpenAIGatewayServiceWithSettings(t, DefaultOpenAIFastPolicySettings())
- account := &Account{Platform: PlatformOpenAI, Type: AccountTypeAPIKey}
- for _, tier := range []string{"auto", "default", "scale"} {
- body := []byte(`{"model":"gpt-5.5","service_tier":"` + tier + `"}`)
- updated, err := svc.applyOpenAIFastPolicyToBody(context.Background(), account, "gpt-5.5", body)
- require.NoError(t, err, "tier %q should pass without error", tier)
- require.Contains(t, string(updated), `"service_tier":"`+tier+`"`,
- "tier %q should be preserved in body under default policy", tier)
- }
- // evaluate 层也应判定为 pass(默认配置没有内置规则)。
- for _, tier := range []string{"auto", "default", "scale"} {
- action, _ := svc.evaluateOpenAIFastPolicy(context.Background(), account, "gpt-5.5", tier)
- require.Equal(t, BetaPolicyActionPass, action, "tier %q should evaluate to pass", tier)
- }
- }
- // TestApplyOpenAIFastPolicyToBody_AllRuleStripsOfficialTiers 验证管理员显式配置
- // ServiceTier=all + Action=filter 规则后,auto/default/scale 等官方 tier 也会
- // 被剥离。这是符合预期的——首条匹配 short-circuit,"all" 覆盖任意已识别 tier。
- func TestApplyOpenAIFastPolicyToBody_AllRuleStripsOfficialTiers(t *testing.T) {
- settings := &OpenAIFastPolicySettings{
- Rules: []OpenAIFastPolicyRule{{
- ServiceTier: OpenAIFastTierAny,
- Action: BetaPolicyActionFilter,
- Scope: BetaPolicyScopeAll,
- }},
- }
- svc := newOpenAIGatewayServiceWithSettings(t, settings)
- account := &Account{Platform: PlatformOpenAI, Type: AccountTypeAPIKey}
- for _, tier := range []string{"auto", "default", "scale", "priority", "flex"} {
- body := []byte(`{"model":"gpt-5.5","service_tier":"` + tier + `"}`)
- updated, err := svc.applyOpenAIFastPolicyToBody(context.Background(), account, "gpt-5.5", body)
- require.NoError(t, err)
- require.NotContains(t, string(updated), `"service_tier"`,
- "tier %q should be stripped under ServiceTier=all + filter rule", tier)
- }
- }
- // TestApplyOpenAIFastPolicyToBody_UnknownTierStripped 验证真未知 tier 仍被剥离
- // (normalize 返回 nil → normalizeResponsesBodyServiceTier 删除字段;
- // applyOpenAIFastPolicyToBody 在 normTier 为空时直接 no-op,因为字段已不可能存在
- // 于经过前置归一化的请求里。这里直接调 apply 验证它对未识别值不会异常)。
- func TestApplyOpenAIFastPolicyToBody_UnknownTierStripped(t *testing.T) {
- svc := newOpenAIGatewayServiceWithSettings(t, DefaultOpenAIFastPolicySettings())
- account := &Account{Platform: PlatformOpenAI, Type: AccountTypeAPIKey}
- // normalize 阶段会将未知值剥离
- require.Nil(t, normalizeOpenAIServiceTier("xxx"))
- // applyOpenAIFastPolicyToBody 收到未识别 tier 时不报错,body 透传不变
- // (不属于本函数职责——上层 normalizeResponsesBodyServiceTier 已剥离)
- body := []byte(`{"model":"gpt-5.5","service_tier":"xxx"}`)
- updated, err := svc.applyOpenAIFastPolicyToBody(context.Background(), account, "gpt-5.5", body)
- require.NoError(t, err)
- require.Equal(t, string(body), string(updated))
- }
- func TestApplyOpenAIFastPolicyToBody_BlockReturnsTypedError(t *testing.T) {
- settings := &OpenAIFastPolicySettings{
- Rules: []OpenAIFastPolicyRule{{
- ServiceTier: OpenAIFastTierPriority,
- Action: BetaPolicyActionBlock,
- Scope: BetaPolicyScopeAll,
- ErrorMessage: "fast mode is blocked for gpt-5.5",
- ModelWhitelist: []string{"gpt-5.5"},
- FallbackAction: BetaPolicyActionPass,
- }},
- }
- svc := newOpenAIGatewayServiceWithSettings(t, settings)
- account := &Account{Platform: PlatformOpenAI, Type: AccountTypeAPIKey}
- body := []byte(`{"model":"gpt-5.5","service_tier":"priority"}`)
- updated, err := svc.applyOpenAIFastPolicyToBody(context.Background(), account, "gpt-5.5", body)
- require.Error(t, err)
- var blocked *OpenAIFastBlockedError
- require.True(t, errors.As(err, &blocked))
- require.Contains(t, blocked.Message, "fast mode is blocked")
- require.Equal(t, string(body), string(updated)) // body not mutated on block
- }
- func TestSetOpenAIFastPolicySettings_Validation(t *testing.T) {
- repo := &openAIFastPolicyRepoStub{values: map[string]string{}}
- svc := NewSettingService(repo, &config.Config{})
- // Invalid action rejected
- err := svc.SetOpenAIFastPolicySettings(context.Background(), &OpenAIFastPolicySettings{
- Rules: []OpenAIFastPolicyRule{{
- ServiceTier: OpenAIFastTierPriority,
- Action: "bogus",
- Scope: BetaPolicyScopeAll,
- }},
- })
- require.Error(t, err)
- // Invalid service_tier rejected
- err = svc.SetOpenAIFastPolicySettings(context.Background(), &OpenAIFastPolicySettings{
- Rules: []OpenAIFastPolicyRule{{
- ServiceTier: "turbo",
- Action: BetaPolicyActionPass,
- Scope: BetaPolicyScopeAll,
- }},
- })
- require.Error(t, err)
- // Non-positive and duplicate user IDs are rejected.
- err = svc.SetOpenAIFastPolicySettings(context.Background(), &OpenAIFastPolicySettings{
- Rules: []OpenAIFastPolicyRule{{
- ServiceTier: OpenAIFastTierPriority,
- Action: BetaPolicyActionPass,
- Scope: BetaPolicyScopeAll,
- UserIDs: []int64{0},
- }},
- })
- require.Error(t, err)
- err = svc.SetOpenAIFastPolicySettings(context.Background(), &OpenAIFastPolicySettings{
- Rules: []OpenAIFastPolicyRule{{
- ServiceTier: OpenAIFastTierPriority,
- Action: BetaPolicyActionPass,
- Scope: BetaPolicyScopeAll,
- UserIDs: []int64{42, 42},
- }},
- })
- require.Error(t, err)
- // Valid settings persisted
- err = svc.SetOpenAIFastPolicySettings(context.Background(), &OpenAIFastPolicySettings{
- Rules: []OpenAIFastPolicyRule{
- {
- ServiceTier: OpenAIFastTierPriority,
- Action: OpenAIFastPolicyActionForcePriority,
- Scope: BetaPolicyScopeAll,
- UserIDs: []int64{42, 43},
- },
- {
- ServiceTier: OpenAIFastTierUltrafast,
- Action: BetaPolicyActionPass,
- Scope: BetaPolicyScopeAll,
- },
- },
- })
- require.NoError(t, err)
- got, err := svc.GetOpenAIFastPolicySettings(context.Background())
- require.NoError(t, err)
- require.Len(t, got.Rules, 2)
- require.Equal(t, OpenAIFastTierPriority, got.Rules[0].ServiceTier)
- require.Equal(t, OpenAIFastPolicyActionForcePriority, got.Rules[0].Action)
- require.Equal(t, []int64{42, 43}, got.Rules[0].UserIDs)
- require.Equal(t, OpenAIFastTierUltrafast, got.Rules[1].ServiceTier)
- require.Equal(t, BetaPolicyActionPass, got.Rules[1].Action)
- }
|