Sub2API v1.0 - AI API 网关（二开初始版本，基于上游 Wei-Shaw/sub2api）

backend/internal/service/account_grok_media_eligibility_test.go

@@ -0,0 +1,180 @@
+//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: "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 TestGrokMediaCapabilityKeepsOnlyUnobservedOAuthAsProbeCandidate(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,
+		}},
+	}
+	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)
+	})
+}