Sub2API v1.0 - AI API 网关（二开初始版本，基于上游 Wei-Shaw/sub2api）
Release / update-version (push) Has been cancelled
Release / build-frontend (push) Has been cancelled
Release / release (push) Has been cancelled
Release / sync-version-file (push) Has been cancelled
CI / shell (push) Canceled after 0s
CI / test (push) Canceled after 0s
CI / frontend (push) Canceled after 0s
CI / golangci-lint (push) Canceled after 0s
Security Scan / backend-security (push) Canceled after 0s
Security Scan / frontend-security (push) Canceled after 0s
Release / update-version (push) Has been cancelled
Release / build-frontend (push) Has been cancelled
Release / release (push) Has been cancelled
Release / sync-version-file (push) Has been cancelled
CI / shell (push) Canceled after 0s
CI / test (push) Canceled after 0s
CI / frontend (push) Canceled after 0s
CI / golangci-lint (push) Canceled after 0s
Security Scan / backend-security (push) Canceled after 0s
Security Scan / frontend-security (push) Canceled after 0s
This commit is contained in:
李建琦
@@ -0,0 +1,419 @@
|
||||
//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)
|
||||
})
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user