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

@@ -5,7 +5,9 @@ import (
"encoding/json"
"fmt"
"io"
"math/rand"
"net/http"
"strconv"
"time"
"github.com/Wei-Shaw/sub2api/internal/pkg/httpclient"
@@ -49,13 +51,38 @@ func (c *driveClient) GetStorageQuota(ctx context.Context, accessToken, proxyURL
return nil, fmt.Errorf("failed to create HTTP client: %w", err)
}
// Retry logic with exponential backoff for rate limits
sleepWithContext := func(d time.Duration) error {
timer := time.NewTimer(d)
defer timer.Stop()
select {
case <-ctx.Done():
return ctx.Err()
case <-timer.C:
return nil
}
}
// Retry logic with exponential backoff (+ jitter) for rate limits and transient failures
var resp *http.Response
maxRetries := 3
rng := rand.New(rand.NewSource(time.Now().UnixNano()))
for attempt := 0; attempt < maxRetries; attempt++ {
if ctx.Err() != nil {
return nil, fmt.Errorf("request cancelled: %w", ctx.Err())
}
resp, err = client.Do(req)
if err != nil {
return nil, fmt.Errorf("failed to execute request: %w", err)
// Network error retry
if attempt < maxRetries-1 {
backoff := time.Duration(1<<uint(attempt)) * time.Second
jitter := time.Duration(rng.Intn(1000)) * time.Millisecond
if err := sleepWithContext(backoff + jitter); err != nil {
return nil, fmt.Errorf("request cancelled: %w", err)
}
continue
}
return nil, fmt.Errorf("network error after %d attempts: %w", maxRetries, err)
}
// Success
@@ -63,18 +90,34 @@ func (c *driveClient) GetStorageQuota(ctx context.Context, accessToken, proxyURL
break
}
// Rate limit - retry with exponential backoff
if resp.StatusCode == http.StatusTooManyRequests && attempt < maxRetries-1 {
// Retry 429, 500, 502, 503 with exponential backoff + jitter
if (resp.StatusCode == http.StatusTooManyRequests ||
resp.StatusCode == http.StatusInternalServerError ||
resp.StatusCode == http.StatusBadGateway ||
resp.StatusCode == http.StatusServiceUnavailable) && attempt < maxRetries-1 {
_ = resp.Body.Close()
backoff := time.Duration(1<<uint(attempt)) * time.Second // 1s, 2s, 4s
time.Sleep(backoff)
backoff := time.Duration(1<<uint(attempt)) * time.Second
jitter := time.Duration(rng.Intn(1000)) * time.Millisecond
if err := sleepWithContext(backoff + jitter); err != nil {
return nil, fmt.Errorf("request cancelled: %w", err)
}
continue
}
// Other errors - return immediately
break
}
if resp == nil {
return nil, fmt.Errorf("request failed: no response received")
}
if resp.StatusCode != http.StatusOK {
body, _ := io.ReadAll(resp.Body)
_ = resp.Body.Close()
return nil, fmt.Errorf("drive API error (status %d): %s", resp.StatusCode, string(body))
// 记录完整错误
fmt.Printf("[DriveClient] API error (status %d): %s\n", resp.StatusCode, string(body))
// 只返回通用错误
return nil, fmt.Errorf("drive API error: status %d", resp.StatusCode)
}
defer func() { _ = resp.Body.Close() }()
@@ -94,10 +137,14 @@ func (c *driveClient) GetStorageQuota(ctx context.Context, accessToken, proxyURL
// Parse limit and usage (handle both string and number formats)
var limit, usage int64
if result.StorageQuota.Limit != "" {
_, _ = fmt.Sscanf(result.StorageQuota.Limit, "%d", &limit)
if val, err := strconv.ParseInt(result.StorageQuota.Limit, 10, 64); err == nil {
limit = val
}
}
if result.StorageQuota.Usage != "" {
_, _ = fmt.Sscanf(result.StorageQuota.Usage, "%d", &usage)
if val, err := strconv.ParseInt(result.StorageQuota.Usage, 10, 64); err == nil {
usage = val
}
}
return &DriveStorageInfo{

View File

@@ -0,0 +1,19 @@
package geminicli
import "testing"
func TestDriveStorageInfo(t *testing.T) {
// 测试 DriveStorageInfo 结构体
info := &DriveStorageInfo{
Limit: 100 * 1024 * 1024 * 1024, // 100GB
Usage: 50 * 1024 * 1024 * 1024, // 50GB
}
if info.Limit != 100*1024*1024*1024 {
t.Errorf("Expected limit 100GB, got %d", info.Limit)
}
if info.Usage != 50*1024*1024*1024 {
t.Errorf("Expected usage 50GB, got %d", info.Usage)
}
}