gemini_oauth_service_test.go 46 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091929394959697989910010110210310410510610710810911011111211311411511611711811912012112212312412512612712812913013113213313413513613713813914014114214314414514614714814915015115215315415515615715815916016116216316416516616716816917017117217317417517617717817918018118218318418518618718818919019119219319419519619719819920020120220320420520620720820921021121221321421521621721821922022122222322422522622722822923023123223323423523623723823924024124224324424524624724824925025125225325425525625725825926026126226326426526626726826927027127227327427527627727827928028128228328428528628728828929029129229329429529629729829930030130230330430530630730830931031131231331431531631731831932032132232332432532632732832933033133233333433533633733833934034134234334434534634734834935035135235335435535635735835936036136236336436536636736836937037137237337437537637737837938038138238338438538638738838939039139239339439539639739839940040140240340440540640740840941041141241341441541641741841942042142242342442542642742842943043143243343443543643743843944044144244344444544644744844945045145245345445545645745845946046146246346446546646746846947047147247347447547647747847948048148248348448548648748848949049149249349449549649749849950050150250350450550650750850951051151251351451551651751851952052152252352452552652752852953053153253353453553653753853954054154254354454554654754854955055155255355455555655755855956056156256356456556656756856957057157257357457557657757857958058158258358458558658758858959059159259359459559659759859960060160260360460560660760860961061161261361461561661761861962062162262362462562662762862963063163263363463563663763863964064164264364464564664764864965065165265365465565665765865966066166266366466566666766866967067167267367467567667767867968068168268368468568668768868969069169269369469569669769869970070170270370470570670770870971071171271371471571671771871972072172272372472572672772872973073173273373473573673773873974074174274374474574674774874975075175275375475575675775875976076176276376476576676776876977077177277377477577677777877978078178278378478578678778878979079179279379479579679779879980080180280380480580680780880981081181281381481581681781881982082182282382482582682782882983083183283383483583683783883984084184284384484584684784884985085185285385485585685785885986086186286386486586686786886987087187287387487587687787887988088188288388488588688788888989089189289389489589689789889990090190290390490590690790890991091191291391491591691791891992092192292392492592692792892993093193293393493593693793893994094194294394494594694794894995095195295395495595695795895996096196296396496596696796896997097197297397497597697797897998098198298398498598698798898999099199299399499599699799899910001001100210031004100510061007100810091010101110121013101410151016101710181019102010211022102310241025102610271028102910301031103210331034103510361037103810391040104110421043104410451046104710481049105010511052105310541055105610571058105910601061106210631064106510661067106810691070107110721073107410751076107710781079108010811082108310841085108610871088108910901091109210931094109510961097109810991100110111021103110411051106110711081109111011111112111311141115111611171118111911201121112211231124112511261127112811291130113111321133113411351136113711381139114011411142114311441145114611471148114911501151115211531154115511561157115811591160116111621163116411651166116711681169117011711172117311741175117611771178117911801181118211831184118511861187118811891190119111921193119411951196119711981199120012011202120312041205120612071208120912101211121212131214121512161217121812191220122112221223122412251226122712281229123012311232123312341235123612371238123912401241124212431244124512461247124812491250125112521253125412551256125712581259126012611262126312641265126612671268126912701271127212731274127512761277127812791280128112821283128412851286128712881289129012911292129312941295129612971298129913001301130213031304130513061307130813091310131113121313131413151316131713181319132013211322132313241325132613271328132913301331133213331334133513361337133813391340134113421343134413451346134713481349135013511352135313541355135613571358135913601361136213631364136513661367136813691370137113721373137413751376137713781379138013811382138313841385138613871388138913901391139213931394139513961397139813991400140114021403140414051406140714081409141014111412141314141415141614171418141914201421142214231424142514261427142814291430143114321433143414351436143714381439144014411442144314441445144614471448144914501451145214531454145514561457145814591460146114621463146414651466146714681469147014711472147314741475147614771478147914801481148214831484148514861487
  1. //go:build unit
  2. package service
  3. import (
  4. "context"
  5. "fmt"
  6. "net/url"
  7. "strings"
  8. "testing"
  9. "time"
  10. "github.com/Wei-Shaw/sub2api/internal/config"
  11. "github.com/Wei-Shaw/sub2api/internal/pkg/geminicli"
  12. "github.com/Wei-Shaw/sub2api/internal/pkg/pagination"
  13. )
  14. // =====================
  15. // 保留原有测试
  16. // =====================
  17. func TestGeminiOAuthService_GenerateAuthURL_RedirectURIStrategy(t *testing.T) {
  18. // NOTE: This test sets process env; it must not run in parallel.
  19. // The built-in Gemini CLI client secret is not embedded in this repository.
  20. // Tests set a dummy secret via env to simulate operator-provided configuration.
  21. t.Setenv(geminicli.GeminiCLIOAuthClientSecretEnv, "test-built-in-secret")
  22. type testCase struct {
  23. name string
  24. cfg *config.Config
  25. oauthType string
  26. projectID string
  27. wantClientID string
  28. wantRedirect string
  29. wantScope string
  30. wantProjectID string
  31. wantErrSubstr string
  32. }
  33. tests := []testCase{
  34. {
  35. name: "google_one uses built-in client when not configured and redirects to upstream",
  36. cfg: &config.Config{
  37. Gemini: config.GeminiConfig{
  38. OAuth: config.GeminiOAuthConfig{},
  39. },
  40. },
  41. oauthType: "google_one",
  42. wantClientID: geminicli.GeminiCLIOAuthClientID,
  43. wantRedirect: geminicli.GeminiCLIRedirectURI,
  44. wantScope: geminicli.DefaultCodeAssistScopes,
  45. wantProjectID: "",
  46. },
  47. {
  48. name: "google_one always forces built-in client even when custom client configured",
  49. cfg: &config.Config{
  50. Gemini: config.GeminiConfig{
  51. OAuth: config.GeminiOAuthConfig{
  52. ClientID: "custom-client-id",
  53. ClientSecret: "custom-client-secret",
  54. },
  55. },
  56. },
  57. oauthType: "google_one",
  58. wantClientID: geminicli.GeminiCLIOAuthClientID,
  59. wantRedirect: geminicli.GeminiCLIRedirectURI,
  60. wantScope: geminicli.DefaultCodeAssistScopes,
  61. wantProjectID: "",
  62. },
  63. {
  64. name: "code_assist always forces built-in client even when custom client configured",
  65. cfg: &config.Config{
  66. Gemini: config.GeminiConfig{
  67. OAuth: config.GeminiOAuthConfig{
  68. ClientID: "custom-client-id",
  69. ClientSecret: "custom-client-secret",
  70. },
  71. },
  72. },
  73. oauthType: "code_assist",
  74. projectID: "my-gcp-project",
  75. wantClientID: geminicli.GeminiCLIOAuthClientID,
  76. wantRedirect: geminicli.GeminiCLIRedirectURI,
  77. wantScope: geminicli.DefaultCodeAssistScopes,
  78. wantProjectID: "my-gcp-project",
  79. },
  80. {
  81. name: "ai_studio requires custom client",
  82. cfg: &config.Config{
  83. Gemini: config.GeminiConfig{
  84. OAuth: config.GeminiOAuthConfig{},
  85. },
  86. },
  87. oauthType: "ai_studio",
  88. wantErrSubstr: "AI Studio OAuth requires a custom OAuth Client",
  89. },
  90. }
  91. for _, tt := range tests {
  92. tt := tt
  93. t.Run(tt.name, func(t *testing.T) {
  94. t.Parallel()
  95. svc := NewGeminiOAuthService(nil, nil, nil, nil, tt.cfg)
  96. got, err := svc.GenerateAuthURL(context.Background(), nil, "https://example.com/auth/callback", tt.projectID, tt.oauthType, "")
  97. if tt.wantErrSubstr != "" {
  98. if err == nil {
  99. t.Fatalf("expected error containing %q, got nil", tt.wantErrSubstr)
  100. }
  101. if !strings.Contains(err.Error(), tt.wantErrSubstr) {
  102. t.Fatalf("expected error containing %q, got: %v", tt.wantErrSubstr, err)
  103. }
  104. return
  105. }
  106. if err != nil {
  107. t.Fatalf("GenerateAuthURL returned error: %v", err)
  108. }
  109. parsed, err := url.Parse(got.AuthURL)
  110. if err != nil {
  111. t.Fatalf("failed to parse auth_url: %v", err)
  112. }
  113. q := parsed.Query()
  114. if gotState := q.Get("state"); gotState != got.State {
  115. t.Fatalf("state mismatch: query=%q result=%q", gotState, got.State)
  116. }
  117. if gotClientID := q.Get("client_id"); gotClientID != tt.wantClientID {
  118. t.Fatalf("client_id mismatch: got=%q want=%q", gotClientID, tt.wantClientID)
  119. }
  120. if gotRedirect := q.Get("redirect_uri"); gotRedirect != tt.wantRedirect {
  121. t.Fatalf("redirect_uri mismatch: got=%q want=%q", gotRedirect, tt.wantRedirect)
  122. }
  123. if gotScope := q.Get("scope"); gotScope != tt.wantScope {
  124. t.Fatalf("scope mismatch: got=%q want=%q", gotScope, tt.wantScope)
  125. }
  126. if gotProjectID := q.Get("project_id"); gotProjectID != tt.wantProjectID {
  127. t.Fatalf("project_id mismatch: got=%q want=%q", gotProjectID, tt.wantProjectID)
  128. }
  129. })
  130. }
  131. }
  132. // =====================
  133. // 新增测试:validateTierID
  134. // =====================
  135. func TestValidateTierID(t *testing.T) {
  136. t.Parallel()
  137. tests := []struct {
  138. name string
  139. tierID string
  140. wantErr bool
  141. }{
  142. {name: "空字符串合法", tierID: "", wantErr: false},
  143. {name: "正常 tier_id", tierID: "google_one_free", wantErr: false},
  144. {name: "包含斜杠", tierID: "tier/sub", wantErr: false},
  145. {name: "包含连字符", tierID: "gcp-standard", wantErr: false},
  146. {name: "纯数字", tierID: "12345", wantErr: false},
  147. {name: "超长字符串(65个字符)", tierID: strings.Repeat("a", 65), wantErr: true},
  148. {name: "刚好64个字符", tierID: strings.Repeat("b", 64), wantErr: false},
  149. {name: "非法字符_空格", tierID: "tier id", wantErr: true},
  150. {name: "非法字符_中文", tierID: "tier_中文", wantErr: true},
  151. {name: "非法字符_特殊符号", tierID: "tier@id", wantErr: true},
  152. {name: "非法字符_感叹号", tierID: "tier!id", wantErr: true},
  153. }
  154. for _, tt := range tests {
  155. tt := tt
  156. t.Run(tt.name, func(t *testing.T) {
  157. t.Parallel()
  158. err := validateTierID(tt.tierID)
  159. if tt.wantErr && err == nil {
  160. t.Fatalf("期望返回错误,但返回 nil")
  161. }
  162. if !tt.wantErr && err != nil {
  163. t.Fatalf("不期望返回错误,但返回: %v", err)
  164. }
  165. })
  166. }
  167. }
  168. // =====================
  169. // 新增测试:canonicalGeminiTierID
  170. // =====================
  171. func TestCanonicalGeminiTierID(t *testing.T) {
  172. t.Parallel()
  173. tests := []struct {
  174. name string
  175. raw string
  176. want string
  177. }{
  178. // 空值
  179. {name: "空字符串", raw: "", want: ""},
  180. {name: "纯空白", raw: " ", want: ""},
  181. // 已规范化的值(直接返回)
  182. {name: "google_one_free", raw: "google_one_free", want: GeminiTierGoogleOneFree},
  183. {name: "google_ai_pro", raw: "google_ai_pro", want: GeminiTierGoogleAIPro},
  184. {name: "google_ai_ultra", raw: "google_ai_ultra", want: GeminiTierGoogleAIUltra},
  185. {name: "gcp_standard", raw: "gcp_standard", want: GeminiTierGCPStandard},
  186. {name: "gcp_enterprise", raw: "gcp_enterprise", want: GeminiTierGCPEnterprise},
  187. {name: "aistudio_free", raw: "aistudio_free", want: GeminiTierAIStudioFree},
  188. {name: "aistudio_paid", raw: "aistudio_paid", want: GeminiTierAIStudioPaid},
  189. {name: "google_one_unknown", raw: "google_one_unknown", want: GeminiTierGoogleOneUnknown},
  190. // 大小写不敏感
  191. {name: "Google_One_Free 大写", raw: "Google_One_Free", want: GeminiTierGoogleOneFree},
  192. {name: "GCP_STANDARD 全大写", raw: "GCP_STANDARD", want: GeminiTierGCPStandard},
  193. // legacy 映射: Google One
  194. {name: "AI_PREMIUM -> google_ai_pro", raw: "AI_PREMIUM", want: GeminiTierGoogleAIPro},
  195. {name: "FREE -> google_one_free", raw: "FREE", want: GeminiTierGoogleOneFree},
  196. {name: "GOOGLE_ONE_BASIC -> google_one_free", raw: "GOOGLE_ONE_BASIC", want: GeminiTierGoogleOneFree},
  197. {name: "GOOGLE_ONE_STANDARD -> google_one_free", raw: "GOOGLE_ONE_STANDARD", want: GeminiTierGoogleOneFree},
  198. {name: "GOOGLE_ONE_UNLIMITED -> google_ai_ultra", raw: "GOOGLE_ONE_UNLIMITED", want: GeminiTierGoogleAIUltra},
  199. {name: "GOOGLE_ONE_UNKNOWN -> google_one_unknown", raw: "GOOGLE_ONE_UNKNOWN", want: GeminiTierGoogleOneUnknown},
  200. // legacy 映射: Code Assist
  201. {name: "STANDARD -> gcp_standard", raw: "STANDARD", want: GeminiTierGCPStandard},
  202. {name: "PRO -> gcp_standard", raw: "PRO", want: GeminiTierGCPStandard},
  203. {name: "LEGACY -> gcp_standard", raw: "LEGACY", want: GeminiTierGCPStandard},
  204. {name: "ENTERPRISE -> gcp_enterprise", raw: "ENTERPRISE", want: GeminiTierGCPEnterprise},
  205. {name: "ULTRA -> gcp_enterprise", raw: "ULTRA", want: GeminiTierGCPEnterprise},
  206. // kebab-case
  207. {name: "standard-tier -> gcp_standard", raw: "standard-tier", want: GeminiTierGCPStandard},
  208. {name: "pro-tier -> gcp_standard", raw: "pro-tier", want: GeminiTierGCPStandard},
  209. {name: "ultra-tier -> gcp_enterprise", raw: "ultra-tier", want: GeminiTierGCPEnterprise},
  210. // 未知值
  211. {name: "unknown_value -> 空", raw: "unknown_value", want: ""},
  212. {name: "random-text -> 空", raw: "random-text", want: ""},
  213. // 带空白
  214. {name: "带前后空白", raw: " google_one_free ", want: GeminiTierGoogleOneFree},
  215. }
  216. for _, tt := range tests {
  217. tt := tt
  218. t.Run(tt.name, func(t *testing.T) {
  219. t.Parallel()
  220. got := canonicalGeminiTierID(tt.raw)
  221. if got != tt.want {
  222. t.Fatalf("canonicalGeminiTierID(%q) = %q, want %q", tt.raw, got, tt.want)
  223. }
  224. })
  225. }
  226. }
  227. // =====================
  228. // 新增测试:canonicalGeminiTierIDForOAuthType
  229. // =====================
  230. func TestCanonicalGeminiTierIDForOAuthType(t *testing.T) {
  231. t.Parallel()
  232. tests := []struct {
  233. name string
  234. oauthType string
  235. tierID string
  236. want string
  237. }{
  238. // google_one 类型过滤
  239. {name: "google_one + google_one_free", oauthType: "google_one", tierID: "google_one_free", want: GeminiTierGoogleOneFree},
  240. {name: "google_one + google_ai_pro", oauthType: "google_one", tierID: "google_ai_pro", want: GeminiTierGoogleAIPro},
  241. {name: "google_one + google_ai_ultra", oauthType: "google_one", tierID: "google_ai_ultra", want: GeminiTierGoogleAIUltra},
  242. {name: "google_one + gcp_standard 被过滤", oauthType: "google_one", tierID: "gcp_standard", want: ""},
  243. {name: "google_one + aistudio_free 被过滤", oauthType: "google_one", tierID: "aistudio_free", want: ""},
  244. {name: "google_one + AI_PREMIUM 遗留映射", oauthType: "google_one", tierID: "AI_PREMIUM", want: GeminiTierGoogleAIPro},
  245. // code_assist 类型过滤
  246. {name: "code_assist + gcp_standard", oauthType: "code_assist", tierID: "gcp_standard", want: GeminiTierGCPStandard},
  247. {name: "code_assist + gcp_enterprise", oauthType: "code_assist", tierID: "gcp_enterprise", want: GeminiTierGCPEnterprise},
  248. {name: "code_assist + google_one_free 被过滤", oauthType: "code_assist", tierID: "google_one_free", want: ""},
  249. {name: "code_assist + aistudio_free 被过滤", oauthType: "code_assist", tierID: "aistudio_free", want: ""},
  250. {name: "code_assist + STANDARD 遗留映射", oauthType: "code_assist", tierID: "STANDARD", want: GeminiTierGCPStandard},
  251. {name: "code_assist + standard-tier kebab", oauthType: "code_assist", tierID: "standard-tier", want: GeminiTierGCPStandard},
  252. // ai_studio 类型过滤
  253. {name: "ai_studio + aistudio_free", oauthType: "ai_studio", tierID: "aistudio_free", want: GeminiTierAIStudioFree},
  254. {name: "ai_studio + aistudio_paid", oauthType: "ai_studio", tierID: "aistudio_paid", want: GeminiTierAIStudioPaid},
  255. {name: "ai_studio + gcp_standard 被过滤", oauthType: "ai_studio", tierID: "gcp_standard", want: ""},
  256. {name: "ai_studio + google_one_free 被过滤", oauthType: "ai_studio", tierID: "google_one_free", want: ""},
  257. // 空值
  258. {name: "空 tierID", oauthType: "google_one", tierID: "", want: ""},
  259. {name: "空 oauthType + 有效 tierID", oauthType: "", tierID: "gcp_standard", want: GeminiTierGCPStandard},
  260. {name: "未知 oauthType 接受规范化值", oauthType: "unknown_type", tierID: "gcp_standard", want: GeminiTierGCPStandard},
  261. // oauthType 大小写和空白
  262. {name: "GOOGLE_ONE 大写", oauthType: "GOOGLE_ONE", tierID: "google_one_free", want: GeminiTierGoogleOneFree},
  263. {name: "oauthType 带空白", oauthType: " code_assist ", tierID: "gcp_standard", want: GeminiTierGCPStandard},
  264. }
  265. for _, tt := range tests {
  266. tt := tt
  267. t.Run(tt.name, func(t *testing.T) {
  268. t.Parallel()
  269. got := canonicalGeminiTierIDForOAuthType(tt.oauthType, tt.tierID)
  270. if got != tt.want {
  271. t.Fatalf("canonicalGeminiTierIDForOAuthType(%q, %q) = %q, want %q", tt.oauthType, tt.tierID, got, tt.want)
  272. }
  273. })
  274. }
  275. }
  276. // =====================
  277. // 新增测试:extractTierIDFromAllowedTiers
  278. // =====================
  279. func TestExtractTierIDFromAllowedTiers(t *testing.T) {
  280. t.Parallel()
  281. tests := []struct {
  282. name string
  283. allowedTiers []geminicli.AllowedTier
  284. want string
  285. }{
  286. {
  287. name: "nil 列表返回 LEGACY",
  288. allowedTiers: nil,
  289. want: "LEGACY",
  290. },
  291. {
  292. name: "空列表返回 LEGACY",
  293. allowedTiers: []geminicli.AllowedTier{},
  294. want: "LEGACY",
  295. },
  296. {
  297. name: "有 IsDefault 的 tier",
  298. allowedTiers: []geminicli.AllowedTier{
  299. {ID: "STANDARD", IsDefault: false},
  300. {ID: "PRO", IsDefault: true},
  301. {ID: "ENTERPRISE", IsDefault: false},
  302. },
  303. want: "PRO",
  304. },
  305. {
  306. name: "没有 IsDefault 取第一个非空",
  307. allowedTiers: []geminicli.AllowedTier{
  308. {ID: "STANDARD", IsDefault: false},
  309. {ID: "ENTERPRISE", IsDefault: false},
  310. },
  311. want: "STANDARD",
  312. },
  313. {
  314. name: "IsDefault 的 ID 为空,取第一个非空",
  315. allowedTiers: []geminicli.AllowedTier{
  316. {ID: "", IsDefault: true},
  317. {ID: "PRO", IsDefault: false},
  318. },
  319. want: "PRO",
  320. },
  321. {
  322. name: "所有 ID 都为空返回 LEGACY",
  323. allowedTiers: []geminicli.AllowedTier{
  324. {ID: "", IsDefault: false},
  325. {ID: " ", IsDefault: false},
  326. },
  327. want: "LEGACY",
  328. },
  329. {
  330. name: "ID 带空白会被 trim",
  331. allowedTiers: []geminicli.AllowedTier{
  332. {ID: " STANDARD ", IsDefault: true},
  333. },
  334. want: "STANDARD",
  335. },
  336. {
  337. name: "单个 tier 且 IsDefault",
  338. allowedTiers: []geminicli.AllowedTier{
  339. {ID: "ENTERPRISE", IsDefault: true},
  340. },
  341. want: "ENTERPRISE",
  342. },
  343. {
  344. name: "单个 tier 非 IsDefault",
  345. allowedTiers: []geminicli.AllowedTier{
  346. {ID: "ENTERPRISE", IsDefault: false},
  347. },
  348. want: "ENTERPRISE",
  349. },
  350. }
  351. for _, tt := range tests {
  352. tt := tt
  353. t.Run(tt.name, func(t *testing.T) {
  354. t.Parallel()
  355. got := extractTierIDFromAllowedTiers(tt.allowedTiers)
  356. if got != tt.want {
  357. t.Fatalf("extractTierIDFromAllowedTiers() = %q, want %q", got, tt.want)
  358. }
  359. })
  360. }
  361. }
  362. // =====================
  363. // 新增测试:inferGoogleOneTier
  364. // =====================
  365. func TestInferGoogleOneTier(t *testing.T) {
  366. t.Parallel()
  367. tests := []struct {
  368. name string
  369. storageBytes int64
  370. want string
  371. }{
  372. // 边界:<= 0
  373. {name: "0 bytes -> unknown", storageBytes: 0, want: GeminiTierGoogleOneUnknown},
  374. {name: "负数 -> unknown", storageBytes: -1, want: GeminiTierGoogleOneUnknown},
  375. // > 100TB -> ultra
  376. {name: "> 100TB -> ultra", storageBytes: int64(StorageTierUnlimited) + 1, want: GeminiTierGoogleAIUltra},
  377. {name: "200TB -> ultra", storageBytes: 200 * int64(TB), want: GeminiTierGoogleAIUltra},
  378. // >= 2TB -> pro (但 <= 100TB)
  379. {name: "正好 2TB -> pro", storageBytes: int64(StorageTierAIPremium), want: GeminiTierGoogleAIPro},
  380. {name: "5TB -> pro", storageBytes: 5 * int64(TB), want: GeminiTierGoogleAIPro},
  381. {name: "100TB 正好 -> pro (不是 > 100TB)", storageBytes: int64(StorageTierUnlimited), want: GeminiTierGoogleAIPro},
  382. // >= 15GB -> free (但 < 2TB)
  383. {name: "正好 15GB -> free", storageBytes: int64(StorageTierFree), want: GeminiTierGoogleOneFree},
  384. {name: "100GB -> free", storageBytes: 100 * int64(GB), want: GeminiTierGoogleOneFree},
  385. {name: "略低于 2TB -> free", storageBytes: int64(StorageTierAIPremium) - 1, want: GeminiTierGoogleOneFree},
  386. // < 15GB -> unknown
  387. {name: "1GB -> unknown", storageBytes: int64(GB), want: GeminiTierGoogleOneUnknown},
  388. {name: "略低于 15GB -> unknown", storageBytes: int64(StorageTierFree) - 1, want: GeminiTierGoogleOneUnknown},
  389. {name: "1 byte -> unknown", storageBytes: 1, want: GeminiTierGoogleOneUnknown},
  390. }
  391. for _, tt := range tests {
  392. tt := tt
  393. t.Run(tt.name, func(t *testing.T) {
  394. t.Parallel()
  395. got := inferGoogleOneTier(tt.storageBytes)
  396. if got != tt.want {
  397. t.Fatalf("inferGoogleOneTier(%d) = %q, want %q", tt.storageBytes, got, tt.want)
  398. }
  399. })
  400. }
  401. }
  402. // =====================
  403. // 新增测试:isNonRetryableGeminiOAuthError
  404. // =====================
  405. func TestIsNonRetryableGeminiOAuthError(t *testing.T) {
  406. t.Parallel()
  407. tests := []struct {
  408. name string
  409. err error
  410. want bool
  411. }{
  412. {name: "invalid_grant", err: fmt.Errorf("error: invalid_grant"), want: true},
  413. {name: "invalid_client", err: fmt.Errorf("oauth error: invalid_client"), want: true},
  414. {name: "unauthorized_client", err: fmt.Errorf("unauthorized_client: mismatch"), want: true},
  415. {name: "access_denied", err: fmt.Errorf("access_denied by user"), want: true},
  416. {name: "普通网络错误", err: fmt.Errorf("connection timeout"), want: false},
  417. {name: "HTTP 500 错误", err: fmt.Errorf("server error 500"), want: false},
  418. {name: "空错误信息", err: fmt.Errorf(""), want: false},
  419. {name: "包含 invalid 但不是完整匹配", err: fmt.Errorf("invalid request"), want: false},
  420. }
  421. for _, tt := range tests {
  422. tt := tt
  423. t.Run(tt.name, func(t *testing.T) {
  424. t.Parallel()
  425. got := isNonRetryableGeminiOAuthError(tt.err)
  426. if got != tt.want {
  427. t.Fatalf("isNonRetryableGeminiOAuthError(%v) = %v, want %v", tt.err, got, tt.want)
  428. }
  429. })
  430. }
  431. }
  432. // =====================
  433. // 新增测试:BuildAccountCredentials
  434. // =====================
  435. func TestGeminiOAuthService_BuildAccountCredentials(t *testing.T) {
  436. t.Parallel()
  437. svc := NewGeminiOAuthService(nil, nil, nil, nil, &config.Config{})
  438. defer svc.Stop()
  439. t.Run("完整字段", func(t *testing.T) {
  440. t.Parallel()
  441. tokenInfo := &GeminiTokenInfo{
  442. AccessToken: "access-123",
  443. RefreshToken: "refresh-456",
  444. ExpiresIn: 3600,
  445. ExpiresAt: 1700000000,
  446. TokenType: "Bearer",
  447. Scope: "openid email",
  448. ProjectID: "my-project",
  449. TierID: "gcp_standard",
  450. OAuthType: "code_assist",
  451. Extra: map[string]any{
  452. "drive_storage_limit": int64(2199023255552),
  453. },
  454. }
  455. creds := svc.BuildAccountCredentials(tokenInfo)
  456. assertCredStr(t, creds, "access_token", "access-123")
  457. assertCredStr(t, creds, "refresh_token", "refresh-456")
  458. assertCredStr(t, creds, "token_type", "Bearer")
  459. assertCredStr(t, creds, "scope", "openid email")
  460. assertCredStr(t, creds, "project_id", "my-project")
  461. assertCredStr(t, creds, "tier_id", "gcp_standard")
  462. assertCredStr(t, creds, "oauth_type", "code_assist")
  463. assertCredStr(t, creds, "expires_at", "1700000000")
  464. if _, ok := creds["drive_storage_limit"]; !ok {
  465. t.Fatal("extra 字段 drive_storage_limit 未包含在 creds 中")
  466. }
  467. })
  468. t.Run("最小字段(仅 access_token 和 expires_at)", func(t *testing.T) {
  469. t.Parallel()
  470. tokenInfo := &GeminiTokenInfo{
  471. AccessToken: "token-only",
  472. ExpiresAt: 1700000000,
  473. }
  474. creds := svc.BuildAccountCredentials(tokenInfo)
  475. assertCredStr(t, creds, "access_token", "token-only")
  476. assertCredStr(t, creds, "expires_at", "1700000000")
  477. // 可选字段不应存在
  478. for _, key := range []string{"refresh_token", "token_type", "scope", "project_id", "tier_id", "oauth_type"} {
  479. if _, ok := creds[key]; ok {
  480. t.Fatalf("不应包含空字段 %q", key)
  481. }
  482. }
  483. })
  484. t.Run("无效 tier_id 被静默跳过", func(t *testing.T) {
  485. t.Parallel()
  486. tokenInfo := &GeminiTokenInfo{
  487. AccessToken: "token",
  488. ExpiresAt: 1700000000,
  489. TierID: "tier with spaces",
  490. }
  491. creds := svc.BuildAccountCredentials(tokenInfo)
  492. if _, ok := creds["tier_id"]; ok {
  493. t.Fatal("无效 tier_id 不应被存入 creds")
  494. }
  495. })
  496. t.Run("超长 tier_id 被静默跳过", func(t *testing.T) {
  497. t.Parallel()
  498. tokenInfo := &GeminiTokenInfo{
  499. AccessToken: "token",
  500. ExpiresAt: 1700000000,
  501. TierID: strings.Repeat("x", 65),
  502. }
  503. creds := svc.BuildAccountCredentials(tokenInfo)
  504. if _, ok := creds["tier_id"]; ok {
  505. t.Fatal("超长 tier_id 不应被存入 creds")
  506. }
  507. })
  508. t.Run("无 extra 字段", func(t *testing.T) {
  509. t.Parallel()
  510. tokenInfo := &GeminiTokenInfo{
  511. AccessToken: "token",
  512. ExpiresAt: 1700000000,
  513. RefreshToken: "rt",
  514. }
  515. creds := svc.BuildAccountCredentials(tokenInfo)
  516. // 仅包含基础字段
  517. if len(creds) != 3 { // access_token, expires_at, refresh_token
  518. t.Fatalf("creds 字段数量不匹配: got=%d want=3, keys=%v", len(creds), credKeys(creds))
  519. }
  520. })
  521. }
  522. // =====================
  523. // 新增测试:GetOAuthConfig
  524. // =====================
  525. func TestGeminiOAuthService_GetOAuthConfig(t *testing.T) {
  526. t.Parallel()
  527. tests := []struct {
  528. name string
  529. cfg *config.Config
  530. wantEnabled bool
  531. }{
  532. {
  533. name: "无自定义 OAuth 客户端",
  534. cfg: &config.Config{
  535. Gemini: config.GeminiConfig{
  536. OAuth: config.GeminiOAuthConfig{},
  537. },
  538. },
  539. wantEnabled: false,
  540. },
  541. {
  542. name: "仅 ClientID 无 ClientSecret",
  543. cfg: &config.Config{
  544. Gemini: config.GeminiConfig{
  545. OAuth: config.GeminiOAuthConfig{
  546. ClientID: "custom-id",
  547. },
  548. },
  549. },
  550. wantEnabled: false,
  551. },
  552. {
  553. name: "仅 ClientSecret 无 ClientID",
  554. cfg: &config.Config{
  555. Gemini: config.GeminiConfig{
  556. OAuth: config.GeminiOAuthConfig{
  557. ClientSecret: "custom-secret",
  558. },
  559. },
  560. },
  561. wantEnabled: false,
  562. },
  563. {
  564. name: "使用内置 Gemini CLI ClientID(不算自定义)",
  565. cfg: &config.Config{
  566. Gemini: config.GeminiConfig{
  567. OAuth: config.GeminiOAuthConfig{
  568. ClientID: geminicli.GeminiCLIOAuthClientID,
  569. ClientSecret: "some-secret",
  570. },
  571. },
  572. },
  573. wantEnabled: false,
  574. },
  575. {
  576. name: "自定义 OAuth 客户端(非内置 ID)",
  577. cfg: &config.Config{
  578. Gemini: config.GeminiConfig{
  579. OAuth: config.GeminiOAuthConfig{
  580. ClientID: "my-custom-client-id",
  581. ClientSecret: "my-custom-client-secret",
  582. },
  583. },
  584. },
  585. wantEnabled: true,
  586. },
  587. {
  588. name: "带空白的自定义客户端",
  589. cfg: &config.Config{
  590. Gemini: config.GeminiConfig{
  591. OAuth: config.GeminiOAuthConfig{
  592. ClientID: " my-custom-client-id ",
  593. ClientSecret: " my-custom-client-secret ",
  594. },
  595. },
  596. },
  597. wantEnabled: true,
  598. },
  599. {
  600. name: "纯空白字符串不算配置",
  601. cfg: &config.Config{
  602. Gemini: config.GeminiConfig{
  603. OAuth: config.GeminiOAuthConfig{
  604. ClientID: " ",
  605. ClientSecret: " ",
  606. },
  607. },
  608. },
  609. wantEnabled: false,
  610. },
  611. }
  612. for _, tt := range tests {
  613. tt := tt
  614. t.Run(tt.name, func(t *testing.T) {
  615. t.Parallel()
  616. svc := NewGeminiOAuthService(nil, nil, nil, nil, tt.cfg)
  617. defer svc.Stop()
  618. result := svc.GetOAuthConfig()
  619. if result.AIStudioOAuthEnabled != tt.wantEnabled {
  620. t.Fatalf("AIStudioOAuthEnabled = %v, want %v", result.AIStudioOAuthEnabled, tt.wantEnabled)
  621. }
  622. // RequiredRedirectURIs 始终包含 AI Studio redirect URI
  623. if len(result.RequiredRedirectURIs) != 1 || result.RequiredRedirectURIs[0] != geminicli.AIStudioOAuthRedirectURI {
  624. t.Fatalf("RequiredRedirectURIs 不匹配: got=%v", result.RequiredRedirectURIs)
  625. }
  626. })
  627. }
  628. }
  629. // =====================
  630. // 新增测试:GeminiOAuthService.Stop
  631. // =====================
  632. func TestGeminiOAuthService_Stop_NoPanic(t *testing.T) {
  633. t.Parallel()
  634. svc := NewGeminiOAuthService(nil, nil, nil, nil, &config.Config{})
  635. // 调用 Stop 不应 panic
  636. svc.Stop()
  637. // 多次调用也不应 panic
  638. svc.Stop()
  639. }
  640. // =====================
  641. // mock: GeminiOAuthClient
  642. // =====================
  643. type mockGeminiOAuthClient struct {
  644. exchangeCodeFunc func(ctx context.Context, oauthType, code, codeVerifier, redirectURI, proxyURL string) (*geminicli.TokenResponse, error)
  645. refreshTokenFunc func(ctx context.Context, oauthType, refreshToken, proxyURL string) (*geminicli.TokenResponse, error)
  646. }
  647. func (m *mockGeminiOAuthClient) ExchangeCode(ctx context.Context, oauthType, code, codeVerifier, redirectURI, proxyURL string) (*geminicli.TokenResponse, error) {
  648. if m.exchangeCodeFunc != nil {
  649. return m.exchangeCodeFunc(ctx, oauthType, code, codeVerifier, redirectURI, proxyURL)
  650. }
  651. panic("ExchangeCode not implemented")
  652. }
  653. func (m *mockGeminiOAuthClient) RefreshToken(ctx context.Context, oauthType, refreshToken, proxyURL string) (*geminicli.TokenResponse, error) {
  654. if m.refreshTokenFunc != nil {
  655. return m.refreshTokenFunc(ctx, oauthType, refreshToken, proxyURL)
  656. }
  657. panic("RefreshToken not implemented")
  658. }
  659. // =====================
  660. // mock: GeminiCliCodeAssistClient
  661. // =====================
  662. type mockGeminiCodeAssistClient struct {
  663. loadCodeAssistFunc func(ctx context.Context, accessToken, proxyURL string, req *geminicli.LoadCodeAssistRequest) (*geminicli.LoadCodeAssistResponse, error)
  664. onboardUserFunc func(ctx context.Context, accessToken, proxyURL string, req *geminicli.OnboardUserRequest) (*geminicli.OnboardUserResponse, error)
  665. }
  666. func (m *mockGeminiCodeAssistClient) LoadCodeAssist(ctx context.Context, accessToken, proxyURL string, req *geminicli.LoadCodeAssistRequest) (*geminicli.LoadCodeAssistResponse, error) {
  667. if m.loadCodeAssistFunc != nil {
  668. return m.loadCodeAssistFunc(ctx, accessToken, proxyURL, req)
  669. }
  670. panic("LoadCodeAssist not implemented")
  671. }
  672. func (m *mockGeminiCodeAssistClient) OnboardUser(ctx context.Context, accessToken, proxyURL string, req *geminicli.OnboardUserRequest) (*geminicli.OnboardUserResponse, error) {
  673. if m.onboardUserFunc != nil {
  674. return m.onboardUserFunc(ctx, accessToken, proxyURL, req)
  675. }
  676. panic("OnboardUser not implemented")
  677. }
  678. // =====================
  679. // mock: ProxyRepository (最小实现)
  680. // =====================
  681. type mockGeminiProxyRepo struct {
  682. getByIDFunc func(ctx context.Context, id int64) (*Proxy, error)
  683. }
  684. func (m *mockGeminiProxyRepo) Create(ctx context.Context, proxy *Proxy) error { panic("not impl") }
  685. func (m *mockGeminiProxyRepo) GetByID(ctx context.Context, id int64) (*Proxy, error) {
  686. if m.getByIDFunc != nil {
  687. return m.getByIDFunc(ctx, id)
  688. }
  689. return nil, fmt.Errorf("proxy not found")
  690. }
  691. func (m *mockGeminiProxyRepo) ListByIDs(ctx context.Context, ids []int64) ([]Proxy, error) {
  692. panic("not impl")
  693. }
  694. func (m *mockGeminiProxyRepo) Update(ctx context.Context, proxy *Proxy) error { panic("not impl") }
  695. func (m *mockGeminiProxyRepo) Delete(ctx context.Context, id int64) error { panic("not impl") }
  696. func (m *mockGeminiProxyRepo) List(ctx context.Context, params pagination.PaginationParams) ([]Proxy, *pagination.PaginationResult, error) {
  697. panic("not impl")
  698. }
  699. func (m *mockGeminiProxyRepo) ListWithFilters(ctx context.Context, params pagination.PaginationParams, protocol, status, search string) ([]Proxy, *pagination.PaginationResult, error) {
  700. panic("not impl")
  701. }
  702. func (m *mockGeminiProxyRepo) ListWithFiltersAndAccountCount(ctx context.Context, params pagination.PaginationParams, protocol, status, search string) ([]ProxyWithAccountCount, *pagination.PaginationResult, error) {
  703. panic("not impl")
  704. }
  705. func (m *mockGeminiProxyRepo) ListActive(ctx context.Context) ([]Proxy, error) { panic("not impl") }
  706. func (m *mockGeminiProxyRepo) ListActiveWithAccountCount(ctx context.Context) ([]ProxyWithAccountCount, error) {
  707. panic("not impl")
  708. }
  709. func (m *mockGeminiProxyRepo) ExistsByHostPortAuth(ctx context.Context, host string, port int, username, password string) (bool, error) {
  710. panic("not impl")
  711. }
  712. func (m *mockGeminiProxyRepo) CountAccountsByProxyID(ctx context.Context, proxyID int64) (int64, error) {
  713. panic("not impl")
  714. }
  715. func (m *mockGeminiProxyRepo) ListAccountSummariesByProxyID(ctx context.Context, proxyID int64) ([]ProxyAccountSummary, error) {
  716. panic("not impl")
  717. }
  718. func (m *mockGeminiProxyRepo) SweepExpiredProxies(ctx context.Context, now time.Time) (int64, error) {
  719. panic("not impl")
  720. }
  721. func (m *mockGeminiProxyRepo) ListAllForFallback(ctx context.Context) ([]Proxy, error) {
  722. panic("not impl")
  723. }
  724. func (m *mockGeminiProxyRepo) CountExpired(ctx context.Context) (int64, error) {
  725. panic("not impl")
  726. }
  727. func (m *mockGeminiProxyRepo) CountExpiringSoon(ctx context.Context, now time.Time) (int64, error) {
  728. panic("not impl")
  729. }
  730. // mockDriveClient implements geminicli.DriveClient for tests.
  731. type mockDriveClient struct {
  732. getStorageQuotaFunc func(ctx context.Context, accessToken, proxyURL string) (*geminicli.DriveStorageInfo, error)
  733. }
  734. func (m *mockDriveClient) GetStorageQuota(ctx context.Context, accessToken, proxyURL string) (*geminicli.DriveStorageInfo, error) {
  735. if m.getStorageQuotaFunc != nil {
  736. return m.getStorageQuotaFunc(ctx, accessToken, proxyURL)
  737. }
  738. return nil, fmt.Errorf("drive API not available in test")
  739. }
  740. // =====================
  741. // 新增测试:GeminiOAuthService.RefreshToken(含重试逻辑)
  742. // =====================
  743. func TestGeminiOAuthService_RefreshToken_Success(t *testing.T) {
  744. t.Parallel()
  745. client := &mockGeminiOAuthClient{
  746. refreshTokenFunc: func(ctx context.Context, oauthType, refreshToken, proxyURL string) (*geminicli.TokenResponse, error) {
  747. return &geminicli.TokenResponse{
  748. AccessToken: "new-access",
  749. RefreshToken: "new-refresh",
  750. TokenType: "Bearer",
  751. ExpiresIn: 3600,
  752. Scope: "openid",
  753. }, nil
  754. },
  755. }
  756. svc := NewGeminiOAuthService(nil, client, nil, nil, &config.Config{})
  757. defer svc.Stop()
  758. info, err := svc.RefreshToken(context.Background(), "code_assist", "old-refresh", "")
  759. if err != nil {
  760. t.Fatalf("RefreshToken 返回错误: %v", err)
  761. }
  762. if info.AccessToken != "new-access" {
  763. t.Fatalf("AccessToken 不匹配: got=%q", info.AccessToken)
  764. }
  765. if info.RefreshToken != "new-refresh" {
  766. t.Fatalf("RefreshToken 不匹配: got=%q", info.RefreshToken)
  767. }
  768. if info.ExpiresAt == 0 {
  769. t.Fatal("ExpiresAt 不应为 0")
  770. }
  771. }
  772. func TestGeminiOAuthService_RefreshToken_NonRetryableError(t *testing.T) {
  773. t.Parallel()
  774. client := &mockGeminiOAuthClient{
  775. refreshTokenFunc: func(ctx context.Context, oauthType, refreshToken, proxyURL string) (*geminicli.TokenResponse, error) {
  776. return nil, fmt.Errorf("invalid_grant: token revoked")
  777. },
  778. }
  779. svc := NewGeminiOAuthService(nil, client, nil, nil, &config.Config{})
  780. defer svc.Stop()
  781. _, err := svc.RefreshToken(context.Background(), "code_assist", "revoked-token", "")
  782. if err == nil {
  783. t.Fatal("RefreshToken 应返回错误(不可重试的 invalid_grant)")
  784. }
  785. if !strings.Contains(err.Error(), "invalid_grant") {
  786. t.Fatalf("错误应包含 invalid_grant: got=%q", err.Error())
  787. }
  788. }
  789. func TestGeminiOAuthService_RefreshToken_RetryableError(t *testing.T) {
  790. t.Parallel()
  791. callCount := 0
  792. client := &mockGeminiOAuthClient{
  793. refreshTokenFunc: func(ctx context.Context, oauthType, refreshToken, proxyURL string) (*geminicli.TokenResponse, error) {
  794. callCount++
  795. if callCount <= 2 {
  796. return nil, fmt.Errorf("temporary network error")
  797. }
  798. return &geminicli.TokenResponse{
  799. AccessToken: "recovered",
  800. ExpiresIn: 3600,
  801. }, nil
  802. },
  803. }
  804. svc := NewGeminiOAuthService(nil, client, nil, nil, &config.Config{})
  805. defer svc.Stop()
  806. info, err := svc.RefreshToken(context.Background(), "code_assist", "rt", "")
  807. if err != nil {
  808. t.Fatalf("RefreshToken 应在重试后成功: %v", err)
  809. }
  810. if info.AccessToken != "recovered" {
  811. t.Fatalf("AccessToken 不匹配: got=%q", info.AccessToken)
  812. }
  813. if callCount < 3 {
  814. t.Fatalf("应至少调用 3 次(2 次失败 + 1 次成功): got=%d", callCount)
  815. }
  816. }
  817. // =====================
  818. // 新增测试:GeminiOAuthService.RefreshAccountToken
  819. // =====================
  820. func TestGeminiOAuthService_RefreshAccountToken_NotGeminiOAuth(t *testing.T) {
  821. t.Parallel()
  822. svc := NewGeminiOAuthService(nil, nil, nil, nil, &config.Config{})
  823. defer svc.Stop()
  824. account := &Account{
  825. Platform: PlatformAnthropic,
  826. Type: AccountTypeOAuth,
  827. }
  828. _, err := svc.RefreshAccountToken(context.Background(), account)
  829. if err == nil {
  830. t.Fatal("应返回错误(非 Gemini OAuth 账号)")
  831. }
  832. if !strings.Contains(err.Error(), "not a Gemini OAuth account") {
  833. t.Fatalf("错误信息不匹配: got=%q", err.Error())
  834. }
  835. }
  836. func TestGeminiOAuthService_RefreshAccountToken_NoRefreshToken(t *testing.T) {
  837. t.Parallel()
  838. svc := NewGeminiOAuthService(nil, nil, nil, nil, &config.Config{})
  839. defer svc.Stop()
  840. account := &Account{
  841. Platform: PlatformGemini,
  842. Type: AccountTypeOAuth,
  843. Credentials: map[string]any{
  844. "access_token": "at",
  845. "oauth_type": "code_assist",
  846. },
  847. }
  848. _, err := svc.RefreshAccountToken(context.Background(), account)
  849. if err == nil {
  850. t.Fatal("应返回错误(无 refresh_token)")
  851. }
  852. if !strings.Contains(err.Error(), "no refresh token") {
  853. t.Fatalf("错误信息不匹配: got=%q", err.Error())
  854. }
  855. }
  856. func TestGeminiOAuthService_RefreshAccountToken_AIStudio(t *testing.T) {
  857. t.Parallel()
  858. client := &mockGeminiOAuthClient{
  859. refreshTokenFunc: func(ctx context.Context, oauthType, refreshToken, proxyURL string) (*geminicli.TokenResponse, error) {
  860. return &geminicli.TokenResponse{
  861. AccessToken: "refreshed-at",
  862. RefreshToken: "refreshed-rt",
  863. ExpiresIn: 3600,
  864. TokenType: "Bearer",
  865. }, nil
  866. },
  867. }
  868. svc := NewGeminiOAuthService(&mockGeminiProxyRepo{}, client, nil, nil, &config.Config{})
  869. defer svc.Stop()
  870. account := &Account{
  871. Platform: PlatformGemini,
  872. Type: AccountTypeOAuth,
  873. Credentials: map[string]any{
  874. "access_token": "old-at",
  875. "refresh_token": "old-rt",
  876. "oauth_type": "ai_studio",
  877. "tier_id": "aistudio_free",
  878. },
  879. }
  880. info, err := svc.RefreshAccountToken(context.Background(), account)
  881. if err != nil {
  882. t.Fatalf("RefreshAccountToken 返回错误: %v", err)
  883. }
  884. if info.AccessToken != "refreshed-at" {
  885. t.Fatalf("AccessToken 不匹配: got=%q", info.AccessToken)
  886. }
  887. if info.OAuthType != "ai_studio" {
  888. t.Fatalf("OAuthType 不匹配: got=%q", info.OAuthType)
  889. }
  890. }
  891. func TestGeminiOAuthService_RefreshAccountToken_CodeAssist_WithProjectID(t *testing.T) {
  892. t.Parallel()
  893. client := &mockGeminiOAuthClient{
  894. refreshTokenFunc: func(ctx context.Context, oauthType, refreshToken, proxyURL string) (*geminicli.TokenResponse, error) {
  895. return &geminicli.TokenResponse{
  896. AccessToken: "refreshed",
  897. RefreshToken: "new-rt",
  898. ExpiresIn: 3600,
  899. }, nil
  900. },
  901. }
  902. svc := NewGeminiOAuthService(&mockGeminiProxyRepo{}, client, nil, nil, &config.Config{})
  903. defer svc.Stop()
  904. account := &Account{
  905. Platform: PlatformGemini,
  906. Type: AccountTypeOAuth,
  907. Credentials: map[string]any{
  908. "access_token": "old-at",
  909. "refresh_token": "old-rt",
  910. "oauth_type": "code_assist",
  911. "project_id": "my-project",
  912. "tier_id": "gcp_standard",
  913. },
  914. }
  915. info, err := svc.RefreshAccountToken(context.Background(), account)
  916. if err != nil {
  917. t.Fatalf("RefreshAccountToken 返回错误: %v", err)
  918. }
  919. if info.ProjectID != "my-project" {
  920. t.Fatalf("ProjectID 应保留: got=%q", info.ProjectID)
  921. }
  922. if info.TierID != GeminiTierGCPStandard {
  923. t.Fatalf("TierID 不匹配: got=%q want=%q", info.TierID, GeminiTierGCPStandard)
  924. }
  925. if info.OAuthType != "code_assist" {
  926. t.Fatalf("OAuthType 不匹配: got=%q", info.OAuthType)
  927. }
  928. }
  929. func TestGeminiOAuthService_RefreshAccountToken_DefaultOAuthType(t *testing.T) {
  930. t.Parallel()
  931. client := &mockGeminiOAuthClient{
  932. refreshTokenFunc: func(ctx context.Context, oauthType, refreshToken, proxyURL string) (*geminicli.TokenResponse, error) {
  933. if oauthType != "code_assist" {
  934. t.Errorf("默认 oauthType 应为 code_assist: got=%q", oauthType)
  935. }
  936. return &geminicli.TokenResponse{
  937. AccessToken: "refreshed",
  938. ExpiresIn: 3600,
  939. }, nil
  940. },
  941. }
  942. svc := NewGeminiOAuthService(&mockGeminiProxyRepo{}, client, nil, nil, &config.Config{})
  943. defer svc.Stop()
  944. // 无 oauth_type 凭据的旧账号
  945. account := &Account{
  946. Platform: PlatformGemini,
  947. Type: AccountTypeOAuth,
  948. Credentials: map[string]any{
  949. "refresh_token": "old-rt",
  950. "project_id": "proj",
  951. "tier_id": "STANDARD",
  952. },
  953. }
  954. info, err := svc.RefreshAccountToken(context.Background(), account)
  955. if err != nil {
  956. t.Fatalf("RefreshAccountToken 返回错误: %v", err)
  957. }
  958. if info.OAuthType != "code_assist" {
  959. t.Fatalf("OAuthType 应默认为 code_assist: got=%q", info.OAuthType)
  960. }
  961. }
  962. func TestGeminiOAuthService_RefreshAccountToken_WithProxy(t *testing.T) {
  963. t.Parallel()
  964. proxyRepo := &mockGeminiProxyRepo{
  965. getByIDFunc: func(ctx context.Context, id int64) (*Proxy, error) {
  966. return &Proxy{
  967. Protocol: "http",
  968. Host: "proxy.test",
  969. Port: 3128,
  970. }, nil
  971. },
  972. }
  973. client := &mockGeminiOAuthClient{
  974. refreshTokenFunc: func(ctx context.Context, oauthType, refreshToken, proxyURL string) (*geminicli.TokenResponse, error) {
  975. if proxyURL != "http://proxy.test:3128" {
  976. t.Errorf("proxyURL 不匹配: got=%q", proxyURL)
  977. }
  978. return &geminicli.TokenResponse{
  979. AccessToken: "refreshed",
  980. ExpiresIn: 3600,
  981. }, nil
  982. },
  983. }
  984. svc := NewGeminiOAuthService(proxyRepo, client, nil, nil, &config.Config{})
  985. defer svc.Stop()
  986. proxyID := int64(5)
  987. account := &Account{
  988. Platform: PlatformGemini,
  989. Type: AccountTypeOAuth,
  990. ProxyID: &proxyID,
  991. Credentials: map[string]any{
  992. "refresh_token": "rt",
  993. "oauth_type": "code_assist",
  994. "project_id": "proj",
  995. },
  996. }
  997. _, err := svc.RefreshAccountToken(context.Background(), account)
  998. if err != nil {
  999. t.Fatalf("RefreshAccountToken 返回错误: %v", err)
  1000. }
  1001. }
  1002. func TestGeminiOAuthService_RefreshAccountToken_CodeAssist_NoProjectID_AutoDetect(t *testing.T) {
  1003. t.Parallel()
  1004. client := &mockGeminiOAuthClient{
  1005. refreshTokenFunc: func(ctx context.Context, oauthType, refreshToken, proxyURL string) (*geminicli.TokenResponse, error) {
  1006. return &geminicli.TokenResponse{
  1007. AccessToken: "at",
  1008. ExpiresIn: 3600,
  1009. }, nil
  1010. },
  1011. }
  1012. codeAssist := &mockGeminiCodeAssistClient{
  1013. loadCodeAssistFunc: func(ctx context.Context, accessToken, proxyURL string, req *geminicli.LoadCodeAssistRequest) (*geminicli.LoadCodeAssistResponse, error) {
  1014. return &geminicli.LoadCodeAssistResponse{
  1015. CloudAICompanionProject: "auto-project-123",
  1016. CurrentTier: &geminicli.TierInfo{ID: "STANDARD"},
  1017. }, nil
  1018. },
  1019. }
  1020. svc := NewGeminiOAuthService(&mockGeminiProxyRepo{}, client, codeAssist, nil, &config.Config{})
  1021. defer svc.Stop()
  1022. account := &Account{
  1023. Platform: PlatformGemini,
  1024. Type: AccountTypeOAuth,
  1025. Credentials: map[string]any{
  1026. "refresh_token": "rt",
  1027. "oauth_type": "code_assist",
  1028. // 无 project_id,触发 fetchProjectID
  1029. },
  1030. }
  1031. info, err := svc.RefreshAccountToken(context.Background(), account)
  1032. if err != nil {
  1033. t.Fatalf("RefreshAccountToken 返回错误: %v", err)
  1034. }
  1035. if info.ProjectID != "auto-project-123" {
  1036. t.Fatalf("ProjectID 应为自动检测值: got=%q", info.ProjectID)
  1037. }
  1038. if info.TierID != GeminiTierGCPStandard {
  1039. t.Fatalf("TierID 不匹配: got=%q", info.TierID)
  1040. }
  1041. }
  1042. func TestGeminiOAuthService_RefreshAccountToken_CodeAssist_NoProjectID_FailsEmpty(t *testing.T) {
  1043. t.Parallel()
  1044. client := &mockGeminiOAuthClient{
  1045. refreshTokenFunc: func(ctx context.Context, oauthType, refreshToken, proxyURL string) (*geminicli.TokenResponse, error) {
  1046. return &geminicli.TokenResponse{
  1047. AccessToken: "at",
  1048. ExpiresIn: 3600,
  1049. }, nil
  1050. },
  1051. }
  1052. // 返回有 currentTier 但无 cloudaicompanionProject 的响应,
  1053. // 使 fetchProjectID 走"已注册用户"路径(尝试 Cloud Resource Manager -> 失败 -> 返回错误),
  1054. // 避免走 onboardUser 路径(5 次重试 x 2 秒 = 10 秒超时)
  1055. codeAssist := &mockGeminiCodeAssistClient{
  1056. loadCodeAssistFunc: func(ctx context.Context, accessToken, proxyURL string, req *geminicli.LoadCodeAssistRequest) (*geminicli.LoadCodeAssistResponse, error) {
  1057. return &geminicli.LoadCodeAssistResponse{
  1058. CurrentTier: &geminicli.TierInfo{ID: "STANDARD"},
  1059. // 无 CloudAICompanionProject
  1060. }, nil
  1061. },
  1062. }
  1063. svc := NewGeminiOAuthService(&mockGeminiProxyRepo{}, client, codeAssist, nil, &config.Config{})
  1064. defer svc.Stop()
  1065. account := &Account{
  1066. Platform: PlatformGemini,
  1067. Type: AccountTypeOAuth,
  1068. Credentials: map[string]any{
  1069. "refresh_token": "rt",
  1070. "oauth_type": "code_assist",
  1071. },
  1072. }
  1073. _, err := svc.RefreshAccountToken(context.Background(), account)
  1074. if err == nil {
  1075. t.Fatal("应返回错误(无法检测 project_id)")
  1076. }
  1077. if !strings.Contains(err.Error(), "project_id") {
  1078. t.Fatalf("错误信息应包含 project_id: got=%q", err.Error())
  1079. }
  1080. }
  1081. func TestGeminiOAuthService_RefreshAccountToken_GoogleOne_FreshCache(t *testing.T) {
  1082. t.Parallel()
  1083. client := &mockGeminiOAuthClient{
  1084. refreshTokenFunc: func(ctx context.Context, oauthType, refreshToken, proxyURL string) (*geminicli.TokenResponse, error) {
  1085. return &geminicli.TokenResponse{
  1086. AccessToken: "at",
  1087. ExpiresIn: 3600,
  1088. }, nil
  1089. },
  1090. }
  1091. svc := NewGeminiOAuthService(&mockGeminiProxyRepo{}, client, nil, nil, &config.Config{})
  1092. defer svc.Stop()
  1093. account := &Account{
  1094. Platform: PlatformGemini,
  1095. Type: AccountTypeOAuth,
  1096. Credentials: map[string]any{
  1097. "refresh_token": "rt",
  1098. "oauth_type": "google_one",
  1099. "project_id": "proj",
  1100. "tier_id": "google_ai_pro",
  1101. },
  1102. Extra: map[string]any{
  1103. // 缓存刷新时间在 24 小时内
  1104. "drive_tier_updated_at": time.Now().Add(-1 * time.Hour).Format(time.RFC3339),
  1105. },
  1106. }
  1107. info, err := svc.RefreshAccountToken(context.Background(), account)
  1108. if err != nil {
  1109. t.Fatalf("RefreshAccountToken 返回错误: %v", err)
  1110. }
  1111. // 缓存新鲜,应使用已有的 tier_id
  1112. if info.TierID != GeminiTierGoogleAIPro {
  1113. t.Fatalf("TierID 应使用缓存值: got=%q want=%q", info.TierID, GeminiTierGoogleAIPro)
  1114. }
  1115. }
  1116. func TestGeminiOAuthService_RefreshAccountToken_GoogleOne_NoTierID_DefaultsFree(t *testing.T) {
  1117. t.Parallel()
  1118. client := &mockGeminiOAuthClient{
  1119. refreshTokenFunc: func(ctx context.Context, oauthType, refreshToken, proxyURL string) (*geminicli.TokenResponse, error) {
  1120. return &geminicli.TokenResponse{
  1121. AccessToken: "at",
  1122. ExpiresIn: 3600,
  1123. }, nil
  1124. },
  1125. }
  1126. svc := NewGeminiOAuthService(&mockGeminiProxyRepo{}, client, nil, &mockDriveClient{}, &config.Config{})
  1127. defer svc.Stop()
  1128. account := &Account{
  1129. Platform: PlatformGemini,
  1130. Type: AccountTypeOAuth,
  1131. Credentials: map[string]any{
  1132. "refresh_token": "rt",
  1133. "oauth_type": "google_one",
  1134. "project_id": "proj",
  1135. // 无 tier_id
  1136. },
  1137. }
  1138. info, err := svc.RefreshAccountToken(context.Background(), account)
  1139. if err != nil {
  1140. t.Fatalf("RefreshAccountToken 返回错误: %v", err)
  1141. }
  1142. // FetchGoogleOneTier 会被调用但 oauthClient(此处 mock)不实现 Drive API,
  1143. // svc.FetchGoogleOneTier 使用真实 DriveClient 会失败,最终回退到默认值。
  1144. // 由于没有 tier_id 且 FetchGoogleOneTier 失败,应默认为 google_one_free
  1145. if info.TierID != GeminiTierGoogleOneFree {
  1146. t.Fatalf("TierID 应为默认 free: got=%q", info.TierID)
  1147. }
  1148. }
  1149. func TestGeminiOAuthService_RefreshAccountToken_UnauthorizedClient_Fallback(t *testing.T) {
  1150. t.Parallel()
  1151. callCount := 0
  1152. client := &mockGeminiOAuthClient{
  1153. refreshTokenFunc: func(ctx context.Context, oauthType, refreshToken, proxyURL string) (*geminicli.TokenResponse, error) {
  1154. callCount++
  1155. if oauthType == "code_assist" {
  1156. return nil, fmt.Errorf("unauthorized_client: client mismatch")
  1157. }
  1158. // ai_studio 路径成功
  1159. return &geminicli.TokenResponse{
  1160. AccessToken: "recovered",
  1161. ExpiresIn: 3600,
  1162. }, nil
  1163. },
  1164. }
  1165. // 启用自定义 OAuth 客户端以触发 fallback 路径
  1166. cfg := &config.Config{
  1167. Gemini: config.GeminiConfig{
  1168. OAuth: config.GeminiOAuthConfig{
  1169. ClientID: "custom-id",
  1170. ClientSecret: "custom-secret",
  1171. },
  1172. },
  1173. }
  1174. svc := NewGeminiOAuthService(&mockGeminiProxyRepo{}, client, nil, nil, cfg)
  1175. defer svc.Stop()
  1176. account := &Account{
  1177. Platform: PlatformGemini,
  1178. Type: AccountTypeOAuth,
  1179. Credentials: map[string]any{
  1180. "refresh_token": "rt",
  1181. "oauth_type": "code_assist",
  1182. "project_id": "proj",
  1183. "tier_id": "gcp_standard",
  1184. },
  1185. }
  1186. info, err := svc.RefreshAccountToken(context.Background(), account)
  1187. if err != nil {
  1188. t.Fatalf("RefreshAccountToken 应在 fallback 后成功: %v", err)
  1189. }
  1190. if info.AccessToken != "recovered" {
  1191. t.Fatalf("AccessToken 不匹配: got=%q", info.AccessToken)
  1192. }
  1193. }
  1194. func TestGeminiOAuthService_RefreshAccountToken_UnauthorizedClient_NoFallback(t *testing.T) {
  1195. t.Parallel()
  1196. client := &mockGeminiOAuthClient{
  1197. refreshTokenFunc: func(ctx context.Context, oauthType, refreshToken, proxyURL string) (*geminicli.TokenResponse, error) {
  1198. return nil, fmt.Errorf("unauthorized_client: client mismatch")
  1199. },
  1200. }
  1201. // 无自定义 OAuth 客户端,无法 fallback
  1202. svc := NewGeminiOAuthService(&mockGeminiProxyRepo{}, client, nil, nil, &config.Config{})
  1203. defer svc.Stop()
  1204. account := &Account{
  1205. Platform: PlatformGemini,
  1206. Type: AccountTypeOAuth,
  1207. Credentials: map[string]any{
  1208. "refresh_token": "rt",
  1209. "oauth_type": "code_assist",
  1210. "project_id": "proj",
  1211. },
  1212. }
  1213. _, err := svc.RefreshAccountToken(context.Background(), account)
  1214. if err == nil {
  1215. t.Fatal("应返回错误(无 fallback)")
  1216. }
  1217. if !strings.Contains(err.Error(), "OAuth client mismatch") {
  1218. t.Fatalf("错误应包含 OAuth client mismatch: got=%q", err.Error())
  1219. }
  1220. }
  1221. // =====================
  1222. // 新增测试:GeminiOAuthService.ExchangeCode
  1223. // =====================
  1224. func TestGeminiOAuthService_ExchangeCode_SessionNotFound(t *testing.T) {
  1225. t.Parallel()
  1226. svc := NewGeminiOAuthService(nil, nil, nil, nil, &config.Config{})
  1227. defer svc.Stop()
  1228. _, err := svc.ExchangeCode(context.Background(), &GeminiExchangeCodeInput{
  1229. SessionID: "nonexistent",
  1230. State: "some-state",
  1231. Code: "some-code",
  1232. })
  1233. if err == nil {
  1234. t.Fatal("应返回错误(session 不存在)")
  1235. }
  1236. if !strings.Contains(err.Error(), "session not found") {
  1237. t.Fatalf("错误信息不匹配: got=%q", err.Error())
  1238. }
  1239. }
  1240. func TestGeminiOAuthService_ExchangeCode_InvalidState(t *testing.T) {
  1241. t.Parallel()
  1242. svc := NewGeminiOAuthService(nil, nil, nil, nil, &config.Config{})
  1243. defer svc.Stop()
  1244. // 手动创建 session(必须设置 CreatedAt,否则会因 TTL 过期被拒绝)
  1245. svc.sessionStore.Set("test-session", &geminicli.OAuthSession{
  1246. State: "correct-state",
  1247. CodeVerifier: "verifier",
  1248. OAuthType: "ai_studio",
  1249. CreatedAt: time.Now(),
  1250. })
  1251. _, err := svc.ExchangeCode(context.Background(), &GeminiExchangeCodeInput{
  1252. SessionID: "test-session",
  1253. State: "wrong-state",
  1254. Code: "code",
  1255. })
  1256. if err == nil {
  1257. t.Fatal("应返回错误(state 不匹配)")
  1258. }
  1259. if !strings.Contains(err.Error(), "invalid state") {
  1260. t.Fatalf("错误信息不匹配: got=%q", err.Error())
  1261. }
  1262. }
  1263. func TestGeminiOAuthService_ExchangeCode_EmptyState(t *testing.T) {
  1264. t.Parallel()
  1265. svc := NewGeminiOAuthService(nil, nil, nil, nil, &config.Config{})
  1266. defer svc.Stop()
  1267. svc.sessionStore.Set("test-session", &geminicli.OAuthSession{
  1268. State: "correct-state",
  1269. CodeVerifier: "verifier",
  1270. CreatedAt: time.Now(),
  1271. })
  1272. _, err := svc.ExchangeCode(context.Background(), &GeminiExchangeCodeInput{
  1273. SessionID: "test-session",
  1274. State: "",
  1275. Code: "code",
  1276. })
  1277. if err == nil {
  1278. t.Fatal("应返回错误(空 state)")
  1279. }
  1280. }
  1281. // =====================
  1282. // 辅助函数
  1283. // =====================
  1284. func assertCredStr(t *testing.T, creds map[string]any, key, want string) {
  1285. t.Helper()
  1286. raw, ok := creds[key]
  1287. if !ok {
  1288. t.Fatalf("creds 缺少 key=%q", key)
  1289. }
  1290. got, ok := raw.(string)
  1291. if !ok {
  1292. t.Fatalf("creds[%q] 不是 string: %T", key, raw)
  1293. }
  1294. if got != want {
  1295. t.Fatalf("creds[%q] = %q, want %q", key, got, want)
  1296. }
  1297. }
  1298. func credKeys(m map[string]any) []string {
  1299. keys := make([]string, 0, len(m))
  1300. for k := range m {
  1301. keys = append(keys, k)
  1302. }
  1303. return keys
  1304. }