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:
@@ -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 {
|
||||
|
||||
Reference in New Issue
Block a user