主要改动: - 后端:重构 Gemini 配额服务,支持多层级配额策略(GCP Standard/Free, Google One, AI Studio, Code Assist) - 后端:优化 OAuth 服务,增强 tier_id 识别和存储逻辑 - 后端:改进用量统计服务,支持不同平台的配额查询 - 后端:优化限流服务,增加临时解除调度状态管理 - 前端:统一四种授权方式的用量显示格式和徽标样式 - 前端:增强账户配额信息展示,支持多种配额类型 - 前端:改进创建和重新授权模态框的用户体验 - 国际化:完善中英文配额相关文案 - 移除 CHANGELOG.md 文件 测试:所有单元测试通过
131 lines
3.7 KiB
Go
131 lines
3.7 KiB
Go
package service
|
|
|
|
import (
|
|
"context"
|
|
"net/url"
|
|
"strings"
|
|
"testing"
|
|
|
|
"github.com/Wei-Shaw/sub2api/internal/config"
|
|
"github.com/Wei-Shaw/sub2api/internal/pkg/geminicli"
|
|
)
|
|
|
|
func TestGeminiOAuthService_GenerateAuthURL_RedirectURIStrategy(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
type testCase struct {
|
|
name string
|
|
cfg *config.Config
|
|
oauthType string
|
|
projectID string
|
|
wantClientID string
|
|
wantRedirect string
|
|
wantScope string
|
|
wantProjectID string
|
|
wantErrSubstr string
|
|
}
|
|
|
|
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) {
|
|
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)
|
|
}
|
|
})
|
|
}
|
|
}
|