| 123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184 |
- //go:build unit
- package service
- import (
- "context"
- "net/http"
- "testing"
- infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
- "github.com/Wei-Shaw/sub2api/internal/pkg/xai"
- "github.com/stretchr/testify/require"
- )
- func TestGrokMediaGenerationEligibility(t *testing.T) {
- weeklyUsagePercent := 12.5
- forbiddenBilling := &xai.BillingSummary{
- StatusCode: http.StatusForbidden,
- WeeklyStatusCode: http.StatusForbidden,
- MonthlyStatusCode: http.StatusForbidden,
- }
- weeklyAllowance := &xai.BillingSummary{
- PeriodType: "weekly",
- UsagePercent: &weeklyUsagePercent,
- StatusCode: http.StatusOK,
- WeeklyStatusCode: http.StatusOK,
- }
- weeklyForbidden := &xai.BillingSummary{
- StatusCode: http.StatusOK,
- WeeklyStatusCode: http.StatusForbidden,
- MonthlyStatusCode: http.StatusOK,
- }
- monthlyForbidden := &xai.BillingSummary{
- StatusCode: http.StatusOK,
- WeeklyStatusCode: http.StatusOK,
- MonthlyStatusCode: http.StatusForbidden,
- }
- tests := []struct {
- name string
- account *Account
- want bool
- wantReason string
- }{
- {name: "nil account", account: nil, want: false, wantReason: "not_grok"},
- {name: "non grok account", account: &Account{Platform: PlatformOpenAI}, want: false, wantReason: "not_grok"},
- {name: "non oauth grok account stays eligible", account: &Account{Platform: PlatformGrok, Type: AccountTypeAPIKey}, want: true, wantReason: "non_oauth"},
- {name: "unobserved oauth fails closed", account: &Account{Platform: PlatformGrok, Type: AccountTypeOAuth}, want: false, wantReason: "billing_unobserved"},
- {name: "inconclusive successful billing requires paid entitlement evidence", account: &Account{Platform: PlatformGrok, Type: AccountTypeOAuth, Extra: map[string]any{grokBillingExtraKey: &xai.BillingSummary{StatusCode: http.StatusOK, Partial: true}}}, want: false, wantReason: "billing_inconclusive"},
- {name: "weekly paid usage is eligible without inferring from period type", account: &Account{Platform: PlatformGrok, Type: AccountTypeOAuth, Extra: map[string]any{grokBillingExtraKey: weeklyAllowance}}, want: true, wantReason: "eligible"},
- {name: "billing forbidden is rejected", account: &Account{Platform: PlatformGrok, Type: AccountTypeOAuth, Extra: map[string]any{grokBillingExtraKey: forbiddenBilling}}, want: false, wantReason: "billing_forbidden"},
- {name: "weekly billing forbidden is rejected after partial success", account: &Account{Platform: PlatformGrok, Type: AccountTypeOAuth, Extra: map[string]any{grokBillingExtraKey: weeklyForbidden}}, want: false, wantReason: "billing_forbidden"},
- {name: "monthly billing forbidden is rejected after partial success", account: &Account{Platform: PlatformGrok, Type: AccountTypeOAuth, Extra: map[string]any{grokBillingExtraKey: monthlyForbidden}}, want: false, wantReason: "billing_forbidden"},
- {name: "malformed billing observation fails closed", account: &Account{Platform: PlatformGrok, Type: AccountTypeOAuth, Extra: map[string]any{grokBillingExtraKey: make(chan int)}}, want: false, wantReason: "billing_unobserved"},
- {name: "malformed override falls back to observations", account: &Account{Platform: PlatformGrok, Type: AccountTypeOAuth, Extra: map[string]any{GrokMediaEligibleExtraKey: "false", grokBillingExtraKey: weeklyAllowance}}, want: true, wantReason: "eligible"},
- {name: "explicit disable wins", account: &Account{Platform: PlatformGrok, Type: AccountTypeOAuth, Extra: map[string]any{GrokMediaEligibleExtraKey: false}}, want: false, wantReason: "override_disabled"},
- {name: "explicit enable wins over forbidden probe", account: &Account{Platform: PlatformGrok, Type: AccountTypeOAuth, Extra: map[string]any{GrokMediaEligibleExtraKey: true, grokBillingExtraKey: forbiddenBilling}}, want: true, wantReason: "override_enabled"},
- }
- for _, tt := range tests {
- t.Run(tt.name, func(t *testing.T) {
- got, reason := tt.account.GrokMediaGenerationEligibility()
- require.Equal(t, tt.want, got)
- require.Equal(t, tt.wantReason, reason)
- })
- }
- }
- func TestGrokMediaCapabilityProbesUnobservedButRejectsInconclusiveOAuth(t *testing.T) {
- unobserved := &Account{Platform: PlatformGrok, Type: AccountTypeOAuth}
- eligible, reason := unobserved.GrokMediaGenerationEligibility()
- require.False(t, eligible)
- require.Equal(t, "billing_unobserved", reason)
- require.True(t, unobserved.SupportsOpenAIEndpointCapability(OpenAIEndpointCapabilityGrokMediaGeneration))
- inconclusive := &Account{
- Platform: PlatformGrok,
- Type: AccountTypeOAuth,
- Extra: map[string]any{grokBillingExtraKey: &xai.BillingSummary{
- StatusCode: http.StatusOK,
- Partial: true,
- }},
- }
- eligible, reason = inconclusive.GrokMediaGenerationEligibility()
- require.False(t, eligible)
- require.Equal(t, "billing_inconclusive", reason)
- require.False(t, inconclusive.SupportsOpenAIEndpointCapability(OpenAIEndpointCapabilityGrokMediaGeneration))
- }
- func TestGrokMediaCapabilityFiltersOnlyGeneration(t *testing.T) {
- account := &Account{
- ID: 1,
- Platform: PlatformGrok,
- Type: AccountTypeOAuth,
- Status: StatusActive,
- Schedulable: true,
- Concurrency: 1,
- Extra: map[string]any{GrokMediaEligibleExtraKey: false},
- }
- require.True(t, account.SupportsOpenAIEndpointCapability(OpenAIEndpointCapabilityChatCompletions))
- require.False(t, account.SupportsOpenAIEndpointCapability(OpenAIEndpointCapabilityGrokMediaGeneration))
- require.False(t, isOpenAICompatibleAccountEligibleForRequest(
- context.Background(), account, PlatformGrok, "grok-imagine-video", false,
- OpenAIEndpointCapabilityGrokMediaGeneration,
- ))
- }
- func TestNormalizeGrokMediaEligibilityExtra(t *testing.T) {
- t.Run("boolean override is accepted", func(t *testing.T) {
- extra, err := normalizeGrokMediaEligibilityExtra(PlatformGrok, map[string]any{GrokMediaEligibleExtraKey: false})
- require.NoError(t, err)
- require.Equal(t, false, extra[GrokMediaEligibleExtraKey])
- })
- t.Run("null clears override", func(t *testing.T) {
- extra, err := normalizeGrokMediaEligibilityExtra(PlatformGrok, map[string]any{GrokMediaEligibleExtraKey: nil})
- require.NoError(t, err)
- require.NotContains(t, extra, GrokMediaEligibleExtraKey)
- })
- t.Run("malformed override is rejected", func(t *testing.T) {
- _, err := normalizeGrokMediaEligibilityExtra(PlatformGrok, map[string]any{GrokMediaEligibleExtraKey: "false"})
- require.Error(t, err)
- require.Equal(t, http.StatusBadRequest, infraerrors.Code(err))
- })
- t.Run("other platforms ignore provider owned value", func(t *testing.T) {
- extra := map[string]any{GrokMediaEligibleExtraKey: "provider-owned"}
- normalized, err := normalizeGrokMediaEligibilityExtra(PlatformOpenAI, extra)
- require.NoError(t, err)
- require.Equal(t, extra, normalized)
- })
- }
- func TestNormalizeGrokMediaEligibilityUpdateExtra(t *testing.T) {
- account := &Account{Platform: PlatformGrok, Extra: map[string]any{GrokMediaEligibleExtraKey: false}}
- t.Run("omitted override preserves current value", func(t *testing.T) {
- input := &UpdateAccountInput{Extra: map[string]any{"quota_used": float64(1)}}
- normalized, err := normalizeGrokMediaEligibilityUpdateExtra(account, input, map[string]any{"quota_used": float64(1)})
- require.NoError(t, err)
- require.Equal(t, false, normalized[GrokMediaEligibleExtraKey])
- })
- t.Run("null removes current override", func(t *testing.T) {
- input := &UpdateAccountInput{Extra: map[string]any{GrokMediaEligibleExtraKey: nil}}
- normalized, err := normalizeGrokMediaEligibilityUpdateExtra(account, input, map[string]any{GrokMediaEligibleExtraKey: nil})
- require.NoError(t, err)
- require.NotContains(t, normalized, GrokMediaEligibleExtraKey)
- require.Contains(t, input.Extra, GrokMediaEligibleExtraKey)
- })
- t.Run("provided boolean replaces current override", func(t *testing.T) {
- input := &UpdateAccountInput{Extra: map[string]any{GrokMediaEligibleExtraKey: true}}
- normalized, err := normalizeGrokMediaEligibilityUpdateExtra(account, input, map[string]any{GrokMediaEligibleExtraKey: true})
- require.NoError(t, err)
- require.Equal(t, true, normalized[GrokMediaEligibleExtraKey])
- })
- t.Run("malformed override is rejected on update", func(t *testing.T) {
- input := &UpdateAccountInput{Extra: map[string]any{GrokMediaEligibleExtraKey: "false"}}
- _, err := normalizeGrokMediaEligibilityUpdateExtra(account, input, nil)
- require.Error(t, err)
- require.Equal(t, http.StatusBadRequest, infraerrors.Code(err))
- })
- t.Run("non grok update is unchanged", func(t *testing.T) {
- input := &UpdateAccountInput{Extra: map[string]any{GrokMediaEligibleExtraKey: "provider-owned"}}
- normalized := map[string]any{GrokMediaEligibleExtraKey: "provider-owned"}
- got, err := normalizeGrokMediaEligibilityUpdateExtra(&Account{Platform: PlatformOpenAI}, input, normalized)
- require.NoError(t, err)
- require.Equal(t, normalized, got)
- })
- }
|