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:
@@ -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{
|
||||
|
||||
19
backend/internal/pkg/geminicli/drive_client_test.go
Normal file
19
backend/internal/pkg/geminicli/drive_client_test.go
Normal 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)
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user