account_grok_media_eligibility_test.go 8.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184
  1. //go:build unit
  2. package service
  3. import (
  4. "context"
  5. "net/http"
  6. "testing"
  7. infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
  8. "github.com/Wei-Shaw/sub2api/internal/pkg/xai"
  9. "github.com/stretchr/testify/require"
  10. )
  11. func TestGrokMediaGenerationEligibility(t *testing.T) {
  12. weeklyUsagePercent := 12.5
  13. forbiddenBilling := &xai.BillingSummary{
  14. StatusCode: http.StatusForbidden,
  15. WeeklyStatusCode: http.StatusForbidden,
  16. MonthlyStatusCode: http.StatusForbidden,
  17. }
  18. weeklyAllowance := &xai.BillingSummary{
  19. PeriodType: "weekly",
  20. UsagePercent: &weeklyUsagePercent,
  21. StatusCode: http.StatusOK,
  22. WeeklyStatusCode: http.StatusOK,
  23. }
  24. weeklyForbidden := &xai.BillingSummary{
  25. StatusCode: http.StatusOK,
  26. WeeklyStatusCode: http.StatusForbidden,
  27. MonthlyStatusCode: http.StatusOK,
  28. }
  29. monthlyForbidden := &xai.BillingSummary{
  30. StatusCode: http.StatusOK,
  31. WeeklyStatusCode: http.StatusOK,
  32. MonthlyStatusCode: http.StatusForbidden,
  33. }
  34. tests := []struct {
  35. name string
  36. account *Account
  37. want bool
  38. wantReason string
  39. }{
  40. {name: "nil account", account: nil, want: false, wantReason: "not_grok"},
  41. {name: "non grok account", account: &Account{Platform: PlatformOpenAI}, want: false, wantReason: "not_grok"},
  42. {name: "non oauth grok account stays eligible", account: &Account{Platform: PlatformGrok, Type: AccountTypeAPIKey}, want: true, wantReason: "non_oauth"},
  43. {name: "unobserved oauth fails closed", account: &Account{Platform: PlatformGrok, Type: AccountTypeOAuth}, want: false, wantReason: "billing_unobserved"},
  44. {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"},
  45. {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"},
  46. {name: "billing forbidden is rejected", account: &Account{Platform: PlatformGrok, Type: AccountTypeOAuth, Extra: map[string]any{grokBillingExtraKey: forbiddenBilling}}, want: false, wantReason: "billing_forbidden"},
  47. {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"},
  48. {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"},
  49. {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"},
  50. {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"},
  51. {name: "explicit disable wins", account: &Account{Platform: PlatformGrok, Type: AccountTypeOAuth, Extra: map[string]any{GrokMediaEligibleExtraKey: false}}, want: false, wantReason: "override_disabled"},
  52. {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"},
  53. }
  54. for _, tt := range tests {
  55. t.Run(tt.name, func(t *testing.T) {
  56. got, reason := tt.account.GrokMediaGenerationEligibility()
  57. require.Equal(t, tt.want, got)
  58. require.Equal(t, tt.wantReason, reason)
  59. })
  60. }
  61. }
  62. func TestGrokMediaCapabilityProbesUnobservedButRejectsInconclusiveOAuth(t *testing.T) {
  63. unobserved := &Account{Platform: PlatformGrok, Type: AccountTypeOAuth}
  64. eligible, reason := unobserved.GrokMediaGenerationEligibility()
  65. require.False(t, eligible)
  66. require.Equal(t, "billing_unobserved", reason)
  67. require.True(t, unobserved.SupportsOpenAIEndpointCapability(OpenAIEndpointCapabilityGrokMediaGeneration))
  68. inconclusive := &Account{
  69. Platform: PlatformGrok,
  70. Type: AccountTypeOAuth,
  71. Extra: map[string]any{grokBillingExtraKey: &xai.BillingSummary{
  72. StatusCode: http.StatusOK,
  73. Partial: true,
  74. }},
  75. }
  76. eligible, reason = inconclusive.GrokMediaGenerationEligibility()
  77. require.False(t, eligible)
  78. require.Equal(t, "billing_inconclusive", reason)
  79. require.False(t, inconclusive.SupportsOpenAIEndpointCapability(OpenAIEndpointCapabilityGrokMediaGeneration))
  80. }
  81. func TestGrokMediaCapabilityFiltersOnlyGeneration(t *testing.T) {
  82. account := &Account{
  83. ID: 1,
  84. Platform: PlatformGrok,
  85. Type: AccountTypeOAuth,
  86. Status: StatusActive,
  87. Schedulable: true,
  88. Concurrency: 1,
  89. Extra: map[string]any{GrokMediaEligibleExtraKey: false},
  90. }
  91. require.True(t, account.SupportsOpenAIEndpointCapability(OpenAIEndpointCapabilityChatCompletions))
  92. require.False(t, account.SupportsOpenAIEndpointCapability(OpenAIEndpointCapabilityGrokMediaGeneration))
  93. require.False(t, isOpenAICompatibleAccountEligibleForRequest(
  94. context.Background(), account, PlatformGrok, "grok-imagine-video", false,
  95. OpenAIEndpointCapabilityGrokMediaGeneration,
  96. ))
  97. }
  98. func TestNormalizeGrokMediaEligibilityExtra(t *testing.T) {
  99. t.Run("boolean override is accepted", func(t *testing.T) {
  100. extra, err := normalizeGrokMediaEligibilityExtra(PlatformGrok, map[string]any{GrokMediaEligibleExtraKey: false})
  101. require.NoError(t, err)
  102. require.Equal(t, false, extra[GrokMediaEligibleExtraKey])
  103. })
  104. t.Run("null clears override", func(t *testing.T) {
  105. extra, err := normalizeGrokMediaEligibilityExtra(PlatformGrok, map[string]any{GrokMediaEligibleExtraKey: nil})
  106. require.NoError(t, err)
  107. require.NotContains(t, extra, GrokMediaEligibleExtraKey)
  108. })
  109. t.Run("malformed override is rejected", func(t *testing.T) {
  110. _, err := normalizeGrokMediaEligibilityExtra(PlatformGrok, map[string]any{GrokMediaEligibleExtraKey: "false"})
  111. require.Error(t, err)
  112. require.Equal(t, http.StatusBadRequest, infraerrors.Code(err))
  113. })
  114. t.Run("other platforms ignore provider owned value", func(t *testing.T) {
  115. extra := map[string]any{GrokMediaEligibleExtraKey: "provider-owned"}
  116. normalized, err := normalizeGrokMediaEligibilityExtra(PlatformOpenAI, extra)
  117. require.NoError(t, err)
  118. require.Equal(t, extra, normalized)
  119. })
  120. }
  121. func TestNormalizeGrokMediaEligibilityUpdateExtra(t *testing.T) {
  122. account := &Account{Platform: PlatformGrok, Extra: map[string]any{GrokMediaEligibleExtraKey: false}}
  123. t.Run("omitted override preserves current value", func(t *testing.T) {
  124. input := &UpdateAccountInput{Extra: map[string]any{"quota_used": float64(1)}}
  125. normalized, err := normalizeGrokMediaEligibilityUpdateExtra(account, input, map[string]any{"quota_used": float64(1)})
  126. require.NoError(t, err)
  127. require.Equal(t, false, normalized[GrokMediaEligibleExtraKey])
  128. })
  129. t.Run("null removes current override", func(t *testing.T) {
  130. input := &UpdateAccountInput{Extra: map[string]any{GrokMediaEligibleExtraKey: nil}}
  131. normalized, err := normalizeGrokMediaEligibilityUpdateExtra(account, input, map[string]any{GrokMediaEligibleExtraKey: nil})
  132. require.NoError(t, err)
  133. require.NotContains(t, normalized, GrokMediaEligibleExtraKey)
  134. require.Contains(t, input.Extra, GrokMediaEligibleExtraKey)
  135. })
  136. t.Run("provided boolean replaces current override", func(t *testing.T) {
  137. input := &UpdateAccountInput{Extra: map[string]any{GrokMediaEligibleExtraKey: true}}
  138. normalized, err := normalizeGrokMediaEligibilityUpdateExtra(account, input, map[string]any{GrokMediaEligibleExtraKey: true})
  139. require.NoError(t, err)
  140. require.Equal(t, true, normalized[GrokMediaEligibleExtraKey])
  141. })
  142. t.Run("malformed override is rejected on update", func(t *testing.T) {
  143. input := &UpdateAccountInput{Extra: map[string]any{GrokMediaEligibleExtraKey: "false"}}
  144. _, err := normalizeGrokMediaEligibilityUpdateExtra(account, input, nil)
  145. require.Error(t, err)
  146. require.Equal(t, http.StatusBadRequest, infraerrors.Code(err))
  147. })
  148. t.Run("non grok update is unchanged", func(t *testing.T) {
  149. input := &UpdateAccountInput{Extra: map[string]any{GrokMediaEligibleExtraKey: "provider-owned"}}
  150. normalized := map[string]any{GrokMediaEligibleExtraKey: "provider-owned"}
  151. got, err := normalizeGrokMediaEligibilityUpdateExtra(&Account{Platform: PlatformOpenAI}, input, normalized)
  152. require.NoError(t, err)
  153. require.Equal(t, normalized, got)
  154. })
  155. }