mirror of
https://github.com/Wei-Shaw/sub2api.git
synced 2026-10-14 03:52:26 +08:00
feeeei
2588675ce7
- 按 base_url 判定上游：非官方地址的 Gemini API Key 账号视为第三方上游 - 第三方上游与 Vertex 不返回本地模拟配额，隐藏等级与配额说明； 账号列表中第三方上游按普通 Key 账号展示今日统计与额度进度条， Vertex 显示 vertex 徽章与今日统计 - 第三方上游的 429 只认显式重试时长，否则走全局 429 兜底冷却， 不按每日配额冷却到 PST 午夜 - 创建账号：第三方 Base URL 不显示 AI Studio 等级选项，也不写入 tier_id； API 密钥卡片标题与说明不再限定 AI Studio，免费层说明只在官方 Base URL 下显示； OAuth 等级回退下拉只在 OAuth 表单显示
81 lines
3.0 KiB
Go
81 lines
3.0 KiB
Go
//go:build unit
|
|
|
|
package service
|
|
|
|
import (
|
|
"context"
|
|
"testing"
|
|
|
|
"github.com/Wei-Shaw/sub2api/internal/config"
|
|
"github.com/stretchr/testify/require"
|
|
)
|
|
|
|
func TestAccountIsGeminiThirdPartyAPIKey(t *testing.T) {
|
|
tests := []struct {
|
|
name string
|
|
account *Account
|
|
expected bool
|
|
}{
|
|
{"nil_account", nil, false},
|
|
{"empty_base_url_is_official", geminiQuotaAPIKeyAccount(""), false},
|
|
{"official_base_url", geminiQuotaAPIKeyAccount("https://generativelanguage.googleapis.com"), false},
|
|
{"official_base_url_trailing_slash", geminiQuotaAPIKeyAccount("https://generativelanguage.googleapis.com/"), false},
|
|
{"official_host_over_http", geminiQuotaAPIKeyAccount("http://generativelanguage.googleapis.com"), true},
|
|
{"official_host_with_path_prefix", geminiQuotaAPIKeyAccount("https://generativelanguage.googleapis.com/proxy"), true},
|
|
{"relay_base_url", geminiQuotaAPIKeyAccount("https://relay.example.com"), true},
|
|
{"oauth_account", &Account{Platform: PlatformGemini, Type: AccountTypeOAuth, Credentials: map[string]any{"base_url": "https://relay.example.com"}}, false},
|
|
{"vertex_service_account", &Account{Platform: PlatformGemini, Type: AccountTypeServiceAccount}, false},
|
|
{"non_gemini_platform", &Account{Platform: PlatformOpenAI, Type: AccountTypeAPIKey, Credentials: map[string]any{"base_url": "https://relay.example.com"}}, false},
|
|
}
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
require.Equal(t, tt.expected, tt.account.IsGeminiThirdPartyAPIKey())
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestGeminiQuotaService_QuotaForAccountSkipsThirdPartyAndVertex(t *testing.T) {
|
|
svc := NewGeminiQuotaService(&config.Config{}, nil)
|
|
ctx := context.Background()
|
|
|
|
quota, ok := svc.QuotaForAccount(ctx, geminiQuotaAPIKeyAccount(""))
|
|
require.True(t, ok)
|
|
require.Equal(t, int64(50), quota.ProRPD)
|
|
|
|
paid := geminiQuotaAPIKeyAccount("https://generativelanguage.googleapis.com")
|
|
paid.Credentials["tier_id"] = GeminiTierAIStudioPaid
|
|
quota, ok = svc.QuotaForAccount(ctx, paid)
|
|
require.True(t, ok)
|
|
require.Equal(t, int64(-1), quota.ProRPD)
|
|
|
|
relay := geminiQuotaAPIKeyAccount("https://relay.example.com")
|
|
relay.Credentials["tier_id"] = GeminiTierAIStudioFree
|
|
_, ok = svc.QuotaForAccount(ctx, relay)
|
|
require.False(t, ok, "第三方上游不套用 AI Studio 模拟配额")
|
|
|
|
vertex := &Account{
|
|
Platform: PlatformGemini,
|
|
Type: AccountTypeServiceAccount,
|
|
Credentials: map[string]any{"project_id": "my-project", "tier_id": "vertex"},
|
|
}
|
|
_, ok = svc.QuotaForAccount(ctx, vertex)
|
|
require.False(t, ok, "Vertex 按量付费不套用模拟配额")
|
|
|
|
googleOne := &Account{
|
|
Platform: PlatformGemini,
|
|
Type: AccountTypeOAuth,
|
|
Credentials: map[string]any{"oauth_type": "google_one", "tier_id": GeminiTierGoogleOneFree},
|
|
}
|
|
quota, ok = svc.QuotaForAccount(ctx, googleOne)
|
|
require.True(t, ok)
|
|
require.Positive(t, quota.SharedRPD)
|
|
}
|
|
|
|
func geminiQuotaAPIKeyAccount(baseURL string) *Account {
|
|
credentials := map[string]any{}
|
|
if baseURL != "" {
|
|
credentials["base_url"] = baseURL
|
|
}
|
|
return &Account{Platform: PlatformGemini, Type: AccountTypeAPIKey, Credentials: credentials}
|
|
}
|