test(gemini): 添加 Drive API 和 OAuth 服务单元测试

- 新增 drive_client_test.go:Drive API 客户端单元测试
- 新增 gemini_oauth_service_test.go:OAuth 服务单元测试
- 重构 account_handler.go:改进 RefreshTier API 实现
- 优化 drive_client.go:增强错误处理和重试逻辑
- 完善 repository 和 service 层:支持批量 tier 刷新
- 更新迁移文件编号:017 -> 024(避免冲突)
This commit is contained in:
IanShaw027
2026-01-01 15:07:16 +08:00
parent 34bbfb5dd2
commit 48764e15a5
9 changed files with 383 additions and 118 deletions

View File

@@ -26,6 +26,17 @@ const (
TierGoogleOneUnlimited = "GOOGLE_ONE_UNLIMITED"
)
const (
GB = 1024 * 1024 * 1024
TB = 1024 * GB
StorageTierUnlimited = 100 * TB // 100TB
StorageTierAIPremium = 2 * TB // 2TB
StorageTierStandard = 200 * GB // 200GB
StorageTierBasic = 100 * GB // 100GB
StorageTierFree = 15 * GB // 15GB
)
type GeminiOAuthService struct {
sessionStore *geminicli.SessionStore
proxyRepo ProxyRepository
@@ -222,31 +233,21 @@ func inferGoogleOneTier(storageBytes int64) string {
return TierGoogleOneUnknown
}
// Unlimited storage (G Suite legacy)
if storageBytes > 100*1024*1024*1024*1024 { // > 100TB
if storageBytes > StorageTierUnlimited {
return TierGoogleOneUnlimited
}
// AI Premium (2TB+)
if storageBytes >= 2*1024*1024*1024*1024 { // >= 2TB
if storageBytes >= StorageTierAIPremium {
return TierAIPremium
}
// Google One Standard (200GB)
if storageBytes >= 200*1024*1024*1024 { // >= 200GB
if storageBytes >= StorageTierStandard {
return TierGoogleOneStandard
}
// Google One Basic (100GB)
if storageBytes >= 100*1024*1024*1024 { // >= 100GB
if storageBytes >= StorageTierBasic {
return TierGoogleOneBasic
}
// Free (15GB)
if storageBytes >= 15*1024*1024*1024 { // >= 15GB
if storageBytes >= StorageTierFree {
return TierFree
}
return TierGoogleOneUnknown
}
@@ -270,6 +271,60 @@ func (s *GeminiOAuthService) FetchGoogleOneTier(ctx context.Context, accessToken
return tierID, storageInfo, nil
}
// RefreshAccountGoogleOneTier 刷新单个账号的 Google One Tier
func (s *GeminiOAuthService) RefreshAccountGoogleOneTier(
ctx context.Context,
account *Account,
) (tierID string, extra map[string]any, credentials map[string]any, err error) {
if account == nil {
return "", nil, nil, fmt.Errorf("account is nil")
}
// 验证账号类型
oauthType, ok := account.Credentials["oauth_type"].(string)
if !ok || oauthType != "google_one" {
return "", nil, nil, fmt.Errorf("not a google_one OAuth account")
}
// 获取 access_token
accessToken, ok := account.Credentials["access_token"].(string)
if !ok || accessToken == "" {
return "", nil, nil, fmt.Errorf("missing access_token")
}
// 获取 proxy URL
var proxyURL string
if account.ProxyID != nil && account.Proxy != nil {
proxyURL = account.Proxy.URL()
}
// 调用 Drive API
tierID, storageInfo, err := s.FetchGoogleOneTier(ctx, accessToken, proxyURL)
if err != nil {
return "", nil, nil, err
}
// 构建 extra 数据(保留原有 extra 字段)
extra = make(map[string]any)
for k, v := range account.Extra {
extra[k] = v
}
if storageInfo != nil {
extra["drive_storage_limit"] = storageInfo.Limit
extra["drive_storage_usage"] = storageInfo.Usage
extra["drive_tier_updated_at"] = time.Now().Format(time.RFC3339)
}
// 构建 credentials 数据
credentials = make(map[string]any)
for k, v := range account.Credentials {
credentials[k] = v
}
credentials["tier_id"] = tierID
return tierID, extra, credentials, nil
}
func (s *GeminiOAuthService) ExchangeCode(ctx context.Context, input *GeminiExchangeCodeInput) (*GeminiTokenInfo, error) {
session, ok := s.sessionStore.Get(input.SessionID)
if !ok {