feat(gemini): 完善 Gemini OAuth 配额系统和用量显示
主要改动： - 后端：重构 Gemini 配额服务，支持多层级配额策略（GCP Standard/Free, Google One, AI Studio, Code Assist） - 后端：优化 OAuth 服务，增强 tier_id 识别和存储逻辑 - 后端：改进用量统计服务，支持不同平台的配额查询 - 后端：优化限流服务，增加临时解除调度状态管理 - 前端：统一四种授权方式的用量显示格式和徽标样式 - 前端：增强账户配额信息展示，支持多种配额类型 - 前端：改进创建和重新授权模态框的用户体验 - 国际化：完善中英文配额相关文案 - 移除 CHANGELOG.md 文件 测试：所有单元测试通过
This commit is contained in:
IanShaw027
@@ -1,50 +1,129 @@
|
||||
package service
|
||||
|
||||
import "testing"
|
||||
import (
|
||||
"context"
|
||||
"net/url"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
func TestInferGoogleOneTier(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
storageBytes int64
|
||||
expectedTier string
|
||||
}{
|
||||
{"Negative storage", -1, TierGoogleOneUnknown},
|
||||
{"Zero storage", 0, TierGoogleOneUnknown},
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/geminicli"
|
||||
)
|
||||
|
||||
// Free tier boundary (15GB)
|
||||
{"Below free tier", 10 * GB, TierGoogleOneUnknown},
|
||||
{"Just below free tier", StorageTierFree - 1, TierGoogleOneUnknown},
|
||||
{"Free tier (15GB)", StorageTierFree, TierFree},
|
||||
func TestGeminiOAuthService_GenerateAuthURL_RedirectURIStrategy(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
// Basic tier boundary (100GB)
|
||||
{"Between free and basic", 50 * GB, TierFree},
|
||||
{"Just below basic tier", StorageTierBasic - 1, TierFree},
|
||||
{"Basic tier (100GB)", StorageTierBasic, TierGoogleOneBasic},
|
||||
type testCase struct {
|
||||
name string
|
||||
cfg *config.Config
|
||||
oauthType string
|
||||
projectID string
|
||||
wantClientID string
|
||||
wantRedirect string
|
||||
wantScope string
|
||||
wantProjectID string
|
||||
wantErrSubstr string
|
||||
}
|
||||
|
||||
// Standard tier boundary (200GB)
|
||||
{"Between basic and standard", 150 * GB, TierGoogleOneBasic},
|
||||
{"Just below standard tier", StorageTierStandard - 1, TierGoogleOneBasic},
|
||||
{"Standard tier (200GB)", StorageTierStandard, TierGoogleOneStandard},
|
||||
|
||||
// AI Premium tier boundary (2TB)
|
||||
{"Between standard and premium", 1 * TB, TierGoogleOneStandard},
|
||||
{"Just below AI Premium tier", StorageTierAIPremium - 1, TierGoogleOneStandard},
|
||||
{"AI Premium tier (2TB)", StorageTierAIPremium, TierAIPremium},
|
||||
|
||||
// Unlimited tier boundary (> 100TB)
|
||||
{"Between premium and unlimited", 50 * TB, TierAIPremium},
|
||||
{"At unlimited threshold (100TB)", StorageTierUnlimited, TierAIPremium},
|
||||
{"Unlimited tier (100TB+)", StorageTierUnlimited + 1, TierGoogleOneUnlimited},
|
||||
{"Unlimited tier (101TB+)", 101 * TB, TierGoogleOneUnlimited},
|
||||
{"Very large storage", 1000 * TB, TierGoogleOneUnlimited},
|
||||
tests := []testCase{
|
||||
{
|
||||
name: "google_one uses built-in client when not configured and redirects to upstream",
|
||||
cfg: &config.Config{
|
||||
Gemini: config.GeminiConfig{
|
||||
OAuth: config.GeminiOAuthConfig{},
|
||||
},
|
||||
},
|
||||
oauthType: "google_one",
|
||||
wantClientID: geminicli.GeminiCLIOAuthClientID,
|
||||
wantRedirect: geminicli.GeminiCLIRedirectURI,
|
||||
wantScope: geminicli.DefaultCodeAssistScopes,
|
||||
wantProjectID: "",
|
||||
},
|
||||
{
|
||||
name: "google_one uses custom client when configured and redirects to localhost",
|
||||
cfg: &config.Config{
|
||||
Gemini: config.GeminiConfig{
|
||||
OAuth: config.GeminiOAuthConfig{
|
||||
ClientID: "custom-client-id",
|
||||
ClientSecret: "custom-client-secret",
|
||||
},
|
||||
},
|
||||
},
|
||||
oauthType: "google_one",
|
||||
wantClientID: "custom-client-id",
|
||||
wantRedirect: geminicli.AIStudioOAuthRedirectURI,
|
||||
wantScope: geminicli.DefaultGoogleOneScopes,
|
||||
wantProjectID: "",
|
||||
},
|
||||
{
|
||||
name: "code_assist always forces built-in client even when custom client configured",
|
||||
cfg: &config.Config{
|
||||
Gemini: config.GeminiConfig{
|
||||
OAuth: config.GeminiOAuthConfig{
|
||||
ClientID: "custom-client-id",
|
||||
ClientSecret: "custom-client-secret",
|
||||
},
|
||||
},
|
||||
},
|
||||
oauthType: "code_assist",
|
||||
projectID: "my-gcp-project",
|
||||
wantClientID: geminicli.GeminiCLIOAuthClientID,
|
||||
wantRedirect: geminicli.GeminiCLIRedirectURI,
|
||||
wantScope: geminicli.DefaultCodeAssistScopes,
|
||||
wantProjectID: "my-gcp-project",
|
||||
},
|
||||
{
|
||||
name: "ai_studio requires custom client",
|
||||
cfg: &config.Config{
|
||||
Gemini: config.GeminiConfig{
|
||||
OAuth: config.GeminiOAuthConfig{},
|
||||
},
|
||||
},
|
||||
oauthType: "ai_studio",
|
||||
wantErrSubstr: "AI Studio OAuth requires a custom OAuth Client",
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
tt := tt
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
result := inferGoogleOneTier(tt.storageBytes)
|
||||
if result != tt.expectedTier {
|
||||
t.Errorf("inferGoogleOneTier(%d) = %s, want %s",
|
||||
tt.storageBytes, result, tt.expectedTier)
|
||||
t.Parallel()
|
||||
|
||||
svc := NewGeminiOAuthService(nil, nil, nil, tt.cfg)
|
||||
got, err := svc.GenerateAuthURL(context.Background(), nil, "https://example.com/auth/callback", tt.projectID, tt.oauthType, "")
|
||||
if tt.wantErrSubstr != "" {
|
||||
if err == nil {
|
||||
t.Fatalf("expected error containing %q, got nil", tt.wantErrSubstr)
|
||||
}
|
||||
if !strings.Contains(err.Error(), tt.wantErrSubstr) {
|
||||
t.Fatalf("expected error containing %q, got: %v", tt.wantErrSubstr, err)
|
||||
}
|
||||
return
|
||||
}
|
||||
if err != nil {
|
||||
t.Fatalf("GenerateAuthURL returned error: %v", err)
|
||||
}
|
||||
|
||||
parsed, err := url.Parse(got.AuthURL)
|
||||
if err != nil {
|
||||
t.Fatalf("failed to parse auth_url: %v", err)
|
||||
}
|
||||
q := parsed.Query()
|
||||
|
||||
if gotState := q.Get("state"); gotState != got.State {
|
||||
t.Fatalf("state mismatch: query=%q result=%q", gotState, got.State)
|
||||
}
|
||||
if gotClientID := q.Get("client_id"); gotClientID != tt.wantClientID {
|
||||
t.Fatalf("client_id mismatch: got=%q want=%q", gotClientID, tt.wantClientID)
|
||||
}
|
||||
if gotRedirect := q.Get("redirect_uri"); gotRedirect != tt.wantRedirect {
|
||||
t.Fatalf("redirect_uri mismatch: got=%q want=%q", gotRedirect, tt.wantRedirect)
|
||||
}
|
||||
if gotScope := q.Get("scope"); gotScope != tt.wantScope {
|
||||
t.Fatalf("scope mismatch: got=%q want=%q", gotScope, tt.wantScope)
|
||||
}
|
||||
if gotProjectID := q.Get("project_id"); gotProjectID != tt.wantProjectID {
|
||||
t.Fatalf("project_id mismatch: got=%q want=%q", gotProjectID, tt.wantProjectID)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user