sub2api/backend/internal/service/grok_quota_fetcher_test.go

//go:build unit

package service

import (
	"net/http"
	"testing"
	"time"

	"github.com/Wei-Shaw/sub2api/internal/pkg/xai"
	"github.com/stretchr/testify/require"
)

func grokInt64PtrForTest(v int64) *int64 { return &v }
func grokIntPtrForTest(v int) *int       { return &v }

func TestGrokQuotaFetcherBuildUsageInfoUnknownUntilFirstSnapshot(t *testing.T) {
	t.Parallel()

	usage := NewGrokQuotaFetcher().BuildUsageInfo(&Account{Platform: PlatformGrok, Type: AccountTypeOAuth})
	require.Equal(t, "passive", usage.Source)
	require.Equal(t, xai.GrokFreeRolling24hTokenLimit, usage.GrokFreeTokenLimit)
	require.Equal(t, "quota_unknown", usage.ErrorCode)
	require.Contains(t, usage.Error, "unknown until billing is probed")
}

func TestGrokQuotaFetcherDoesNotTreatGrok45ResponsesWindowAsHeavy(t *testing.T) {
	t.Parallel()

	// 8300 / 53M is the grok-4.5 Responses rate-limit window, not a plan fingerprint.
	reqLimit, tokLimit := int64(8300), int64(53_000_000)
	fresh := time.Now().UTC().Format(time.RFC3339)
	account := &Account{
		Platform: PlatformGrok,
		Type:     AccountTypeOAuth,
		Credentials: map[string]any{
			"subscription_tier": "SuperGrokPro",
		},
		Extra: map[string]any{
			grokQuotaSnapshotExtraKey: &xai.QuotaSnapshot{
				Requests:          &xai.QuotaWindow{Limit: &reqLimit},
				Tokens:            &xai.QuotaWindow{Limit: &tokLimit},
				LastHeadersSeenAt: fresh,
				HeadersObserved:   true,
				UpdatedAt:         fresh,
			},
		},
	}

	usage := NewGrokQuotaFetcher().BuildUsageInfo(account)
	require.Equal(t, "supergrok", usage.SubscriptionTier)
}

func TestGrokQuotaFetcherUsesGrok45ResponsesWindowAsHeavy(t *testing.T) {
	t.Parallel()

	reqLimit, tokLimit := int64(8300), int64(53_000_000)
	fresh := time.Now().UTC().Format(time.RFC3339)
	account := &Account{
		Platform: PlatformGrok,
		Type:     AccountTypeOAuth,
		Credentials: map[string]any{
			"subscription_tier": "SuperGrokPro",
		},
		Extra: map[string]any{
			grokQuotaSnapshotExtraKey: &xai.QuotaSnapshot{
				Model:             "grok-4.5",
				Requests:          &xai.QuotaWindow{Limit: &reqLimit},
				Tokens:            &xai.QuotaWindow{Limit: &tokLimit},
				LastHeadersSeenAt: fresh,
				HeadersObserved:   true,
				UpdatedAt:         fresh,
			},
		},
	}

	usage := NewGrokQuotaFetcher().BuildUsageInfo(account)
	require.Equal(t, "supergrok_heavy", usage.SubscriptionTier)
}

func TestGrokQuotaFetcherJWTBeatsAmbiguousSuperGrokProQuota(t *testing.T) {
	t.Parallel()

	heavyReq := int64(8300)
	fresh := time.Now().UTC().Format(time.RFC3339)
	account := &Account{
		Platform: PlatformGrok,
		Type:     AccountTypeOAuth,
		Credentials: map[string]any{
			"access_token":      makeGrokOAuthJWT(map[string]any{"tier": 1}),
			"subscription_tier": "SuperGrokPro",
		},
		Extra: map[string]any{
			grokQuotaSnapshotExtraKey: &xai.QuotaSnapshot{
				Requests:          &xai.QuotaWindow{Limit: &heavyReq},
				LastHeadersSeenAt: fresh,
				HeadersObserved:   true,
				UpdatedAt:         fresh,
			},
		},
	}

	usage := NewGrokQuotaFetcher().BuildUsageInfo(account)
	require.Equal(t, "supergrok", usage.SubscriptionTier)
}

func TestGrokQuotaFetcherPrefersLiveJWTTierOverStaleBillingPlan(t *testing.T) {
	t.Parallel()

	account := &Account{
		Platform: PlatformGrok,
		Type:     AccountTypeOAuth,
		Credentials: map[string]any{
			"access_token":      makeGrokOAuthJWT(map[string]any{"tier": 0}),
			"subscription_tier": "supergrok_heavy",
		},
		Extra: map[string]any{
			grokBillingExtraKey: &xai.BillingSummary{
				Plan:       "SuperGrok Heavy",
				StatusCode: http.StatusOK,
				UpdatedAt:  "2030-01-01T00:00:00Z",
			},
		},
	}

	usage := NewGrokQuotaFetcher().BuildUsageInfo(account)
	require.Equal(t, "free", usage.SubscriptionTier)
	require.Equal(t, "free", usage.SubscriptionTierRaw)
}

func TestGrokQuotaFetcherUsesCredentialTierWhenBillingHasNoPlan(t *testing.T) {
	t.Parallel()

	account := &Account{
		Platform: PlatformGrok,
		Type:     AccountTypeOAuth,
		Credentials: map[string]any{
			"subscription_tier":  " FREE ",
			"entitlement_status": " active ",
		},
		Extra: map[string]any{
			grokBillingExtraKey: &xai.BillingSummary{
				PeriodType: "weekly",
				StatusCode: http.StatusOK,
				UpdatedAt:  "2030-01-01T00:00:00Z",
			},
		},
	}

	usage := NewGrokQuotaFetcher().BuildUsageInfo(account)

	require.NotNil(t, usage.GrokBilling)
	require.Equal(t, "free", usage.SubscriptionTier)
	require.Equal(t, "FREE", usage.SubscriptionTierRaw)
	require.Equal(t, "active", usage.GrokEntitlementStatus)
}

func TestGrokQuotaFetcherBuildUsageInfoFromSnapshot(t *testing.T) {
	t.Parallel()

	updatedAt := "2030-01-01T00:00:00Z"
	account := &Account{
		Platform: PlatformGrok,
		Type:     AccountTypeOAuth,
		Extra: map[string]any{
			grokQuotaSnapshotExtraKey: &xai.QuotaSnapshot{
				Requests: &xai.QuotaWindow{
					Limit:     grokInt64PtrForTest(100),
					Remaining: grokInt64PtrForTest(12),
					ResetAt:   updatedAt,
				},
				Tokens: &xai.QuotaWindow{
					Limit:     grokInt64PtrForTest(1000),
					Remaining: grokInt64PtrForTest(900),
				},
				RetryAfterSeconds: grokIntPtrForTest(30),
				SubscriptionTier:  "supergrok",
				EntitlementStatus: "active",
				StatusCode:        http.StatusTooManyRequests,
				LastProbeAt:       updatedAt,
				LastHeadersSeenAt: updatedAt,
				UpdatedAt:         updatedAt,
			},
		},
	}

	usage := NewGrokQuotaFetcher().BuildUsageInfo(account)
	require.Equal(t, "passive", usage.Source)
	require.Equal(t, "rate_limited", usage.ErrorCode)
	require.Equal(t, "observed", usage.GrokQuotaSnapshotState)
	require.Equal(t, "supergrok", usage.SubscriptionTier)
	require.Equal(t, "active", usage.GrokEntitlementStatus)
	require.Equal(t, int64(100), *usage.GrokRequestQuota.Limit)
	require.Equal(t, int64(12), *usage.GrokRequestQuota.Remaining)
	require.Equal(t, 30, *usage.GrokRetryAfterSeconds)
	require.NotNil(t, usage.UpdatedAt)
	require.Equal(t, updatedAt, usage.GrokLastQuotaProbeAt)
	require.Equal(t, updatedAt, usage.GrokLastHeadersSeenAt)
	require.Equal(t, http.StatusTooManyRequests, usage.GrokLastStatusCode)
	require.True(t, usage.UpdatedAt.Equal(time.Date(2030, 1, 1, 0, 0, 0, 0, time.UTC)))
}

func TestGrokQuotaFetcherSnapshotErrorOverridesSuccessfulBillingStatus(t *testing.T) {
	t.Parallel()

	updatedAt := "2030-01-01T00:00:00Z"
	account := &Account{
		Platform: PlatformGrok,
		Type:     AccountTypeOAuth,
		Extra: map[string]any{
			grokBillingExtraKey: &xai.BillingSummary{
				PeriodType: "weekly",
				StatusCode: http.StatusOK,
				UpdatedAt:  updatedAt,
			},
			grokQuotaSnapshotExtraKey: &xai.QuotaSnapshot{
				StatusCode: http.StatusTooManyRequests,
				UpdatedAt:  updatedAt,
			},
		},
	}

	usage := NewGrokQuotaFetcher().BuildUsageInfo(account)

	require.Equal(t, "rate_limited", usage.ErrorCode)
	require.Equal(t, http.StatusTooManyRequests, usage.GrokLastStatusCode)
}

func TestGrokQuotaFetcherNewerSuccessfulActiveProbeClearsBillingForbidden(t *testing.T) {
	t.Parallel()

	billingAt := "2030-01-01T00:00:00Z"
	probeAt := "2030-01-01T00:05:00Z"
	account := &Account{
		Platform: PlatformGrok,
		Type:     AccountTypeOAuth,
		Credentials: map[string]any{
			"entitlement_status": "forbidden",
		},
		Extra: map[string]any{
			grokBillingExtraKey: &xai.BillingSummary{
				StatusCode: http.StatusForbidden,
				UpdatedAt:  billingAt,
			},
			grokQuotaSnapshotExtraKey: &xai.QuotaSnapshot{
				StatusCode:        http.StatusOK,
				ObservationSource: "active_probe",
				LastProbeAt:       probeAt,
				UpdatedAt:         probeAt,
			},
		},
	}

	usage := NewGrokQuotaFetcher().BuildUsageInfo(account)

	require.False(t, usage.IsForbidden)
	require.Empty(t, usage.ForbiddenType)
	require.Empty(t, usage.ErrorCode)
	require.Empty(t, usage.GrokEntitlementStatus)
	require.Equal(t, http.StatusOK, usage.GrokLastStatusCode)
	require.Equal(t, probeAt, usage.GrokLastQuotaProbeAt)
	require.NotNil(t, usage.UpdatedAt)
	require.True(t, usage.UpdatedAt.Equal(time.Date(2030, 1, 1, 0, 5, 0, 0, time.UTC)))
}

func TestGrokQuotaFetcherSameSecondSuccessfulActiveProbeClearsBillingForbidden(t *testing.T) {
	t.Parallel()

	observedAt := "2030-01-01T00:05:00Z"
	account := &Account{
		Platform: PlatformGrok,
		Type:     AccountTypeOAuth,
		Extra: map[string]any{
			grokBillingExtraKey: &xai.BillingSummary{
				StatusCode: http.StatusForbidden,
				UpdatedAt:  observedAt,
			},
			grokQuotaSnapshotExtraKey: &xai.QuotaSnapshot{
				StatusCode:        http.StatusOK,
				ObservationSource: "active_probe",
				LastProbeAt:       observedAt,
				UpdatedAt:         observedAt,
			},
		},
	}

	usage := NewGrokQuotaFetcher().BuildUsageInfo(account)

	require.False(t, usage.IsForbidden)
	require.Empty(t, usage.ForbiddenType)
	require.Empty(t, usage.ErrorCode)
	require.Equal(t, http.StatusOK, usage.GrokLastStatusCode)
}

func TestGrokQuotaFetcherDoesNotClearBillingForbiddenWithoutNewerSuccessfulActiveProbe(t *testing.T) {
	t.Parallel()

	tests := []struct {
		name     string
		snapshot xai.QuotaSnapshot
	}{
		{
			name: "older active probe",
			snapshot: xai.QuotaSnapshot{
				StatusCode:        http.StatusOK,
				ObservationSource: "active_probe",
				LastProbeAt:       "2030-01-01T00:04:59Z",
				UpdatedAt:         "2030-01-01T00:04:59Z",
			},
		},
		{
			name: "newer passive response",
			snapshot: xai.QuotaSnapshot{
				StatusCode:        http.StatusOK,
				ObservationSource: "upstream_response",
				UpdatedAt:         "2030-01-01T00:05:01Z",
			},
		},
		{
			name: "newer failed active probe",
			snapshot: xai.QuotaSnapshot{
				StatusCode:        http.StatusTooManyRequests,
				ObservationSource: "active_probe",
				LastProbeAt:       "2030-01-01T00:05:01Z",
				UpdatedAt:         "2030-01-01T00:05:01Z",
			},
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			t.Parallel()
			account := &Account{
				Platform: PlatformGrok,
				Type:     AccountTypeOAuth,
				Extra: map[string]any{
					grokBillingExtraKey: &xai.BillingSummary{
						StatusCode: http.StatusForbidden,
						UpdatedAt:  "2030-01-01T00:05:00Z",
					},
					grokQuotaSnapshotExtraKey: tt.snapshot,
				},
			}

			usage := NewGrokQuotaFetcher().BuildUsageInfo(account)

			require.True(t, usage.IsForbidden)
			require.Equal(t, "forbidden", usage.ForbiddenType)
			require.Equal(t, "forbidden", usage.ErrorCode)
		})
	}
}

func TestGrokQuotaFetcherBuildUsageInfoFromNoHeadersProbe(t *testing.T) {
	t.Parallel()

	probedAt := "2030-01-01T00:00:00Z"
	account := &Account{
		Platform: PlatformGrok,
		Type:     AccountTypeOAuth,
		Extra: map[string]any{
			grokQuotaSnapshotExtraKey: xai.QuotaSnapshot{
				StatusCode:        http.StatusOK,
				HeadersObserved:   false,
				ObservationSource: "active_probe",
				LastProbeAt:       probedAt,
				UpdatedAt:         probedAt,
			},
		},
	}

	usage := NewGrokQuotaFetcher().BuildUsageInfo(account)
	require.Equal(t, "quota_unknown", usage.ErrorCode)
	require.Equal(t, "no_headers", usage.GrokQuotaSnapshotState)
	require.Contains(t, usage.Error, "No xAI quota headers observed")
	require.Equal(t, probedAt, usage.GrokLastQuotaProbeAt)
	require.Empty(t, usage.GrokLastHeadersSeenAt)
	require.Equal(t, http.StatusOK, usage.GrokLastStatusCode)
	require.Nil(t, usage.GrokRequestQuota)
	require.Nil(t, usage.GrokTokenQuota)
}

func TestGrokQuotaFetcherClassifiesForbiddenAndReauth(t *testing.T) {
	t.Parallel()

	tests := []struct {
		name        string
		statusCode  int
		wantReauth  bool
		wantForbid  bool
		wantCode    string
		wantEntitle string
	}{
		{name: "reauth", statusCode: http.StatusUnauthorized, wantReauth: true, wantCode: "unauthenticated"},
		{name: "forbidden", statusCode: http.StatusForbidden, wantForbid: true, wantCode: "forbidden", wantEntitle: "forbidden"},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			t.Parallel()
			account := &Account{
				Platform: PlatformGrok,
				Type:     AccountTypeOAuth,
				Extra: map[string]any{
					grokQuotaSnapshotExtraKey: xai.QuotaSnapshot{
						StatusCode:      tt.statusCode,
						HeadersObserved: true,
						UpdatedAt:       time.Now().UTC().Format(time.RFC3339),
					},
				},
			}
			usage := NewGrokQuotaFetcher().BuildUsageInfo(account)
			require.Equal(t, tt.wantReauth, usage.NeedsReauth)
			require.Equal(t, tt.wantForbid, usage.IsForbidden)
			require.Equal(t, tt.wantCode, usage.ErrorCode)
			require.Equal(t, tt.wantEntitle, usage.GrokEntitlementStatus)
		})
	}
}