client_test.go 58 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091929394959697989910010110210310410510610710810911011111211311411511611711811912012112212312412512612712812913013113213313413513613713813914014114214314414514614714814915015115215315415515615715815916016116216316416516616716816917017117217317417517617717817918018118218318418518618718818919019119219319419519619719819920020120220320420520620720820921021121221321421521621721821922022122222322422522622722822923023123223323423523623723823924024124224324424524624724824925025125225325425525625725825926026126226326426526626726826927027127227327427527627727827928028128228328428528628728828929029129229329429529629729829930030130230330430530630730830931031131231331431531631731831932032132232332432532632732832933033133233333433533633733833934034134234334434534634734834935035135235335435535635735835936036136236336436536636736836937037137237337437537637737837938038138238338438538638738838939039139239339439539639739839940040140240340440540640740840941041141241341441541641741841942042142242342442542642742842943043143243343443543643743843944044144244344444544644744844945045145245345445545645745845946046146246346446546646746846947047147247347447547647747847948048148248348448548648748848949049149249349449549649749849950050150250350450550650750850951051151251351451551651751851952052152252352452552652752852953053153253353453553653753853954054154254354454554654754854955055155255355455555655755855956056156256356456556656756856957057157257357457557657757857958058158258358458558658758858959059159259359459559659759859960060160260360460560660760860961061161261361461561661761861962062162262362462562662762862963063163263363463563663763863964064164264364464564664764864965065165265365465565665765865966066166266366466566666766866967067167267367467567667767867968068168268368468568668768868969069169269369469569669769869970070170270370470570670770870971071171271371471571671771871972072172272372472572672772872973073173273373473573673773873974074174274374474574674774874975075175275375475575675775875976076176276376476576676776876977077177277377477577677777877978078178278378478578678778878979079179279379479579679779879980080180280380480580680780880981081181281381481581681781881982082182282382482582682782882983083183283383483583683783883984084184284384484584684784884985085185285385485585685785885986086186286386486586686786886987087187287387487587687787887988088188288388488588688788888989089189289389489589689789889990090190290390490590690790890991091191291391491591691791891992092192292392492592692792892993093193293393493593693793893994094194294394494594694794894995095195295395495595695795895996096196296396496596696796896997097197297397497597697797897998098198298398498598698798898999099199299399499599699799899910001001100210031004100510061007100810091010101110121013101410151016101710181019102010211022102310241025102610271028102910301031103210331034103510361037103810391040104110421043104410451046104710481049105010511052105310541055105610571058105910601061106210631064106510661067106810691070107110721073107410751076107710781079108010811082108310841085108610871088108910901091109210931094109510961097109810991100110111021103110411051106110711081109111011111112111311141115111611171118111911201121112211231124112511261127112811291130113111321133113411351136113711381139114011411142114311441145114611471148114911501151115211531154115511561157115811591160116111621163116411651166116711681169117011711172117311741175117611771178117911801181118211831184118511861187118811891190119111921193119411951196119711981199120012011202120312041205120612071208120912101211121212131214121512161217121812191220122112221223122412251226122712281229123012311232123312341235123612371238123912401241124212431244124512461247124812491250125112521253125412551256125712581259126012611262126312641265126612671268126912701271127212731274127512761277127812791280128112821283128412851286128712881289129012911292129312941295129612971298129913001301130213031304130513061307130813091310131113121313131413151316131713181319132013211322132313241325132613271328132913301331133213331334133513361337133813391340134113421343134413451346134713481349135013511352135313541355135613571358135913601361136213631364136513661367136813691370137113721373137413751376137713781379138013811382138313841385138613871388138913901391139213931394139513961397139813991400140114021403140414051406140714081409141014111412141314141415141614171418141914201421142214231424142514261427142814291430143114321433143414351436143714381439144014411442144314441445144614471448144914501451145214531454145514561457145814591460146114621463146414651466146714681469147014711472147314741475147614771478147914801481148214831484148514861487148814891490149114921493149414951496149714981499150015011502150315041505150615071508150915101511151215131514151515161517151815191520152115221523152415251526152715281529153015311532153315341535153615371538153915401541154215431544154515461547154815491550155115521553155415551556155715581559156015611562156315641565156615671568156915701571157215731574157515761577157815791580158115821583158415851586158715881589159015911592159315941595159615971598159916001601160216031604160516061607160816091610161116121613161416151616161716181619162016211622162316241625162616271628162916301631163216331634163516361637163816391640164116421643164416451646164716481649165016511652165316541655165616571658165916601661166216631664166516661667166816691670167116721673167416751676167716781679168016811682168316841685168616871688168916901691169216931694169516961697169816991700170117021703170417051706170717081709171017111712171317141715171617171718171917201721172217231724172517261727172817291730173117321733173417351736173717381739174017411742174317441745174617471748174917501751175217531754175517561757175817591760176117621763176417651766176717681769177017711772177317741775177617771778177917801781178217831784178517861787178817891790179117921793179417951796179717981799180018011802180318041805180618071808180918101811181218131814181518161817181818191820182118221823182418251826182718281829183018311832183318341835183618371838183918401841184218431844184518461847184818491850185118521853185418551856185718581859
  1. //go:build unit
  2. package antigravity
  3. import (
  4. "context"
  5. "encoding/json"
  6. "fmt"
  7. "net"
  8. "net/http"
  9. "net/http/httptest"
  10. "net/url"
  11. "strings"
  12. "testing"
  13. "time"
  14. )
  15. const defaultFetchAvailableModelsBodyLimit int64 = 8 << 20
  16. // ---------------------------------------------------------------------------
  17. // NewAPIRequestWithURL
  18. // ---------------------------------------------------------------------------
  19. func TestNewAPIRequestWithURL_普通请求(t *testing.T) {
  20. ctx := context.Background()
  21. baseURL := "https://example.com"
  22. action := "generateContent"
  23. token := "test-token"
  24. body := []byte(`{"prompt":"hello"}`)
  25. req, err := NewAPIRequestWithURL(ctx, baseURL, action, token, body)
  26. if err != nil {
  27. t.Fatalf("创建请求失败: %v", err)
  28. }
  29. // 验证 URL 不含 ?alt=sse
  30. expectedURL := "https://example.com/v1internal:generateContent"
  31. if req.URL.String() != expectedURL {
  32. t.Errorf("URL 不匹配: got %s, want %s", req.URL.String(), expectedURL)
  33. }
  34. // 验证请求方法
  35. if req.Method != http.MethodPost {
  36. t.Errorf("请求方法不匹配: got %s, want POST", req.Method)
  37. }
  38. // 验证 Headers
  39. if ct := req.Header.Get("Content-Type"); ct != "application/json" {
  40. t.Errorf("Content-Type 不匹配: got %s", ct)
  41. }
  42. if auth := req.Header.Get("Authorization"); auth != "Bearer test-token" {
  43. t.Errorf("Authorization 不匹配: got %s", auth)
  44. }
  45. if ua := req.Header.Get("User-Agent"); ua != GetUserAgent() {
  46. t.Errorf("User-Agent 不匹配: got %s, want %s", ua, GetUserAgent())
  47. }
  48. }
  49. func TestNewAPIRequestWithURL_流式请求(t *testing.T) {
  50. ctx := context.Background()
  51. baseURL := "https://example.com"
  52. action := "streamGenerateContent"
  53. token := "tok"
  54. body := []byte(`{}`)
  55. req, err := NewAPIRequestWithURL(ctx, baseURL, action, token, body)
  56. if err != nil {
  57. t.Fatalf("创建请求失败: %v", err)
  58. }
  59. expectedURL := "https://example.com/v1internal:streamGenerateContent?alt=sse"
  60. if req.URL.String() != expectedURL {
  61. t.Errorf("URL 不匹配: got %s, want %s", req.URL.String(), expectedURL)
  62. }
  63. }
  64. func TestNewAPIRequestWithURL_空Body(t *testing.T) {
  65. ctx := context.Background()
  66. req, err := NewAPIRequestWithURL(ctx, "https://example.com", "test", "tok", nil)
  67. if err != nil {
  68. t.Fatalf("创建请求失败: %v", err)
  69. }
  70. if req.Body == nil {
  71. t.Error("Body 应该非 nil(bytes.NewReader(nil) 会返回空 reader)")
  72. }
  73. }
  74. // ---------------------------------------------------------------------------
  75. // NewAPIRequest
  76. // ---------------------------------------------------------------------------
  77. func TestNewAPIRequest_使用默认URL(t *testing.T) {
  78. ctx := context.Background()
  79. req, err := NewAPIRequest(ctx, "generateContent", "tok", []byte(`{}`))
  80. if err != nil {
  81. t.Fatalf("创建请求失败: %v", err)
  82. }
  83. expected := BaseURL + "/v1internal:generateContent"
  84. if req.URL.String() != expected {
  85. t.Errorf("URL 不匹配: got %s, want %s", req.URL.String(), expected)
  86. }
  87. }
  88. // ---------------------------------------------------------------------------
  89. // TierInfo.UnmarshalJSON
  90. // ---------------------------------------------------------------------------
  91. func TestTierInfo_UnmarshalJSON_字符串格式(t *testing.T) {
  92. data := []byte(`"free-tier"`)
  93. var tier TierInfo
  94. if err := tier.UnmarshalJSON(data); err != nil {
  95. t.Fatalf("反序列化失败: %v", err)
  96. }
  97. if tier.ID != "free-tier" {
  98. t.Errorf("ID 不匹配: got %s, want free-tier", tier.ID)
  99. }
  100. if tier.Name != "" {
  101. t.Errorf("Name 应为空: got %s", tier.Name)
  102. }
  103. }
  104. func TestTierInfo_UnmarshalJSON_对象格式(t *testing.T) {
  105. data := []byte(`{"id":"g1-pro-tier","name":"Pro","description":"Pro plan"}`)
  106. var tier TierInfo
  107. if err := tier.UnmarshalJSON(data); err != nil {
  108. t.Fatalf("反序列化失败: %v", err)
  109. }
  110. if tier.ID != "g1-pro-tier" {
  111. t.Errorf("ID 不匹配: got %s, want g1-pro-tier", tier.ID)
  112. }
  113. if tier.Name != "Pro" {
  114. t.Errorf("Name 不匹配: got %s, want Pro", tier.Name)
  115. }
  116. if tier.Description != "Pro plan" {
  117. t.Errorf("Description 不匹配: got %s, want Pro plan", tier.Description)
  118. }
  119. }
  120. func TestTierInfo_UnmarshalJSON_null(t *testing.T) {
  121. data := []byte(`null`)
  122. var tier TierInfo
  123. if err := tier.UnmarshalJSON(data); err != nil {
  124. t.Fatalf("反序列化 null 失败: %v", err)
  125. }
  126. if tier.ID != "" {
  127. t.Errorf("null 场景下 ID 应为空: got %s", tier.ID)
  128. }
  129. }
  130. func TestTierInfo_UnmarshalJSON_空数据(t *testing.T) {
  131. data := []byte(``)
  132. var tier TierInfo
  133. if err := tier.UnmarshalJSON(data); err != nil {
  134. t.Fatalf("反序列化空数据失败: %v", err)
  135. }
  136. if tier.ID != "" {
  137. t.Errorf("空数据场景下 ID 应为空: got %s", tier.ID)
  138. }
  139. }
  140. func TestTierInfo_UnmarshalJSON_空格包裹null(t *testing.T) {
  141. data := []byte(` null `)
  142. var tier TierInfo
  143. if err := tier.UnmarshalJSON(data); err != nil {
  144. t.Fatalf("反序列化空格 null 失败: %v", err)
  145. }
  146. if tier.ID != "" {
  147. t.Errorf("空格 null 场景下 ID 应为空: got %s", tier.ID)
  148. }
  149. }
  150. func TestTierInfo_UnmarshalJSON_通过JSON嵌套结构(t *testing.T) {
  151. // 模拟 LoadCodeAssistResponse 中的嵌套反序列化
  152. jsonData := `{"currentTier":"free-tier","paidTier":{"id":"g1-ultra-tier","name":"Ultra"}}`
  153. var resp LoadCodeAssistResponse
  154. if err := json.Unmarshal([]byte(jsonData), &resp); err != nil {
  155. t.Fatalf("反序列化嵌套结构失败: %v", err)
  156. }
  157. if resp.CurrentTier == nil || resp.CurrentTier.ID != "free-tier" {
  158. t.Errorf("CurrentTier 不匹配: got %+v", resp.CurrentTier)
  159. }
  160. if resp.PaidTier == nil || resp.PaidTier.ID != "g1-ultra-tier" {
  161. t.Errorf("PaidTier 不匹配: got %+v", resp.PaidTier)
  162. }
  163. }
  164. // ---------------------------------------------------------------------------
  165. // LoadCodeAssistResponse.GetTier
  166. // ---------------------------------------------------------------------------
  167. func TestGetTier_PaidTier优先(t *testing.T) {
  168. resp := &LoadCodeAssistResponse{
  169. CurrentTier: &TierInfo{ID: "free-tier"},
  170. PaidTier: &PaidTierInfo{ID: "g1-pro-tier"},
  171. }
  172. if got := resp.GetTier(); got != "g1-pro-tier" {
  173. t.Errorf("应返回 paidTier: got %s", got)
  174. }
  175. }
  176. func TestGetTier_回退到CurrentTier(t *testing.T) {
  177. resp := &LoadCodeAssistResponse{
  178. CurrentTier: &TierInfo{ID: "free-tier"},
  179. }
  180. if got := resp.GetTier(); got != "free-tier" {
  181. t.Errorf("应返回 currentTier: got %s", got)
  182. }
  183. }
  184. func TestGetTier_PaidTier为空ID(t *testing.T) {
  185. resp := &LoadCodeAssistResponse{
  186. CurrentTier: &TierInfo{ID: "free-tier"},
  187. PaidTier: &PaidTierInfo{ID: ""},
  188. }
  189. // paidTier.ID 为空时应回退到 currentTier
  190. if got := resp.GetTier(); got != "free-tier" {
  191. t.Errorf("paidTier.ID 为空时应回退到 currentTier: got %s", got)
  192. }
  193. }
  194. func TestGetAvailableCredits(t *testing.T) {
  195. resp := &LoadCodeAssistResponse{
  196. PaidTier: &PaidTierInfo{
  197. ID: "g1-pro-tier",
  198. AvailableCredits: []AvailableCredit{
  199. {
  200. CreditType: "GOOGLE_ONE_AI",
  201. CreditAmount: "25",
  202. MinimumCreditAmountForUsage: "5",
  203. },
  204. },
  205. },
  206. }
  207. credits := resp.GetAvailableCredits()
  208. if len(credits) != 1 {
  209. t.Fatalf("AI Credits 数量不匹配: got %d", len(credits))
  210. }
  211. if credits[0].GetAmount() != 25 {
  212. t.Errorf("CreditAmount 解析不正确: got %v", credits[0].GetAmount())
  213. }
  214. if credits[0].GetMinimumAmount() != 5 {
  215. t.Errorf("MinimumCreditAmountForUsage 解析不正确: got %v", credits[0].GetMinimumAmount())
  216. }
  217. }
  218. func TestGetTier_两者都为nil(t *testing.T) {
  219. resp := &LoadCodeAssistResponse{}
  220. if got := resp.GetTier(); got != "" {
  221. t.Errorf("两者都为 nil 时应返回空字符串: got %s", got)
  222. }
  223. }
  224. func TestTierIDToPlanType(t *testing.T) {
  225. tests := []struct {
  226. tierID string
  227. want string
  228. }{
  229. {"free-tier", "Free"},
  230. {"g1-pro-tier", "Pro"},
  231. {"g1-ultra-tier", "Ultra"},
  232. {"FREE-TIER", "Free"},
  233. {"", "Free"},
  234. {"unknown-tier", "unknown-tier"},
  235. }
  236. for _, tt := range tests {
  237. t.Run(tt.tierID, func(t *testing.T) {
  238. if got := TierIDToPlanType(tt.tierID); got != tt.want {
  239. t.Errorf("TierIDToPlanType(%q) = %q, want %q", tt.tierID, got, tt.want)
  240. }
  241. })
  242. }
  243. }
  244. // ---------------------------------------------------------------------------
  245. // NewClient
  246. // ---------------------------------------------------------------------------
  247. func mustNewClient(t *testing.T, proxyURL string) *Client {
  248. t.Helper()
  249. client, err := NewClient(proxyURL)
  250. if err != nil {
  251. t.Fatalf("NewClient(%q) failed: %v", proxyURL, err)
  252. }
  253. return client
  254. }
  255. func TestNewClient_无代理(t *testing.T) {
  256. client, err := NewClient("")
  257. if err != nil {
  258. t.Fatalf("NewClient 返回错误: %v", err)
  259. }
  260. if client == nil {
  261. t.Fatal("NewClient 返回 nil")
  262. }
  263. if client.httpClient == nil {
  264. t.Fatal("httpClient 为 nil")
  265. }
  266. if client.httpClient.Timeout != clientTimeout {
  267. t.Errorf("Timeout 不匹配: got %v, want %v", client.httpClient.Timeout, clientTimeout)
  268. }
  269. // 无代理时 Transport 应为 nil(使用默认)
  270. if client.httpClient.Transport != nil {
  271. t.Error("无代理时 Transport 应为 nil")
  272. }
  273. }
  274. func TestNewClient_有代理(t *testing.T) {
  275. client, err := NewClient("http://proxy.example.com:8080")
  276. if err != nil {
  277. t.Fatalf("NewClient 返回错误: %v", err)
  278. }
  279. if client == nil {
  280. t.Fatal("NewClient 返回 nil")
  281. }
  282. if client.httpClient.Transport == nil {
  283. t.Fatal("有代理时 Transport 不应为 nil")
  284. }
  285. }
  286. func TestNewClient_空格代理(t *testing.T) {
  287. client, err := NewClient(" ")
  288. if err != nil {
  289. t.Fatalf("NewClient 返回错误: %v", err)
  290. }
  291. if client == nil {
  292. t.Fatal("NewClient 返回 nil")
  293. }
  294. // 空格代理应等同于无代理
  295. if client.httpClient.Transport != nil {
  296. t.Error("空格代理 Transport 应为 nil")
  297. }
  298. }
  299. func TestNewClient_无效代理URL(t *testing.T) {
  300. // 无效 URL 应返回 error
  301. _, err := NewClient("://invalid")
  302. if err == nil {
  303. t.Fatal("无效代理 URL 应返回错误")
  304. }
  305. if !strings.Contains(err.Error(), "invalid proxy URL") {
  306. t.Errorf("错误信息应包含 'invalid proxy URL': got %s", err.Error())
  307. }
  308. }
  309. // ---------------------------------------------------------------------------
  310. // IsConnectionError
  311. // ---------------------------------------------------------------------------
  312. func TestIsConnectionError_nil(t *testing.T) {
  313. if IsConnectionError(nil) {
  314. t.Error("nil 错误不应判定为连接错误")
  315. }
  316. }
  317. func TestIsConnectionError_超时错误(t *testing.T) {
  318. // 使用 net.OpError 包装超时
  319. err := &net.OpError{
  320. Op: "dial",
  321. Net: "tcp",
  322. Err: &timeoutError{},
  323. }
  324. if !IsConnectionError(err) {
  325. t.Error("超时错误应判定为连接错误")
  326. }
  327. }
  328. // timeoutError 实现 net.Error 接口用于测试
  329. type timeoutError struct{}
  330. func (e *timeoutError) Error() string { return "timeout" }
  331. func (e *timeoutError) Timeout() bool { return true }
  332. func (e *timeoutError) Temporary() bool { return true }
  333. func TestIsConnectionError_netOpError(t *testing.T) {
  334. err := &net.OpError{
  335. Op: "dial",
  336. Net: "tcp",
  337. Err: fmt.Errorf("connection refused"),
  338. }
  339. if !IsConnectionError(err) {
  340. t.Error("net.OpError 应判定为连接错误")
  341. }
  342. }
  343. func TestIsConnectionError_urlError(t *testing.T) {
  344. err := &url.Error{
  345. Op: "Get",
  346. URL: "https://example.com",
  347. Err: fmt.Errorf("some error"),
  348. }
  349. if !IsConnectionError(err) {
  350. t.Error("url.Error 应判定为连接错误")
  351. }
  352. }
  353. func TestIsConnectionError_普通错误(t *testing.T) {
  354. err := fmt.Errorf("some random error")
  355. if IsConnectionError(err) {
  356. t.Error("普通错误不应判定为连接错误")
  357. }
  358. }
  359. func TestIsConnectionError_包装的netOpError(t *testing.T) {
  360. inner := &net.OpError{
  361. Op: "dial",
  362. Net: "tcp",
  363. Err: fmt.Errorf("connection refused"),
  364. }
  365. err := fmt.Errorf("wrapping: %w", inner)
  366. if !IsConnectionError(err) {
  367. t.Error("被包装的 net.OpError 应判定为连接错误")
  368. }
  369. }
  370. // ---------------------------------------------------------------------------
  371. // shouldFallbackToNextURL
  372. // ---------------------------------------------------------------------------
  373. func TestShouldFallbackToNextURL_连接错误(t *testing.T) {
  374. err := &net.OpError{Op: "dial", Net: "tcp", Err: fmt.Errorf("refused")}
  375. if !shouldFallbackToNextURL(err, 0) {
  376. t.Error("连接错误应触发 URL 降级")
  377. }
  378. }
  379. func TestShouldFallbackToNextURL_状态码(t *testing.T) {
  380. tests := []struct {
  381. name string
  382. statusCode int
  383. want bool
  384. }{
  385. {"429 Too Many Requests", http.StatusTooManyRequests, true},
  386. {"408 Request Timeout", http.StatusRequestTimeout, true},
  387. {"404 Not Found", http.StatusNotFound, true},
  388. {"500 Internal Server Error", http.StatusInternalServerError, true},
  389. {"502 Bad Gateway", http.StatusBadGateway, true},
  390. {"503 Service Unavailable", http.StatusServiceUnavailable, true},
  391. {"200 OK", http.StatusOK, false},
  392. {"201 Created", http.StatusCreated, false},
  393. {"400 Bad Request", http.StatusBadRequest, false},
  394. {"401 Unauthorized", http.StatusUnauthorized, false},
  395. {"403 Forbidden", http.StatusForbidden, false},
  396. }
  397. for _, tt := range tests {
  398. t.Run(tt.name, func(t *testing.T) {
  399. got := shouldFallbackToNextURL(nil, tt.statusCode)
  400. if got != tt.want {
  401. t.Errorf("shouldFallbackToNextURL(nil, %d) = %v, want %v", tt.statusCode, got, tt.want)
  402. }
  403. })
  404. }
  405. }
  406. func TestShouldFallbackToNextURL_无错误且200(t *testing.T) {
  407. if shouldFallbackToNextURL(nil, http.StatusOK) {
  408. t.Error("无错误且 200 不应触发 URL 降级")
  409. }
  410. }
  411. // ---------------------------------------------------------------------------
  412. // Client.ExchangeCode (使用 httptest)
  413. // ---------------------------------------------------------------------------
  414. func TestClient_ExchangeCode_成功(t *testing.T) {
  415. old := defaultClientSecret
  416. defaultClientSecret = "test-secret"
  417. t.Cleanup(func() { defaultClientSecret = old })
  418. server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
  419. // 验证请求方法
  420. if r.Method != http.MethodPost {
  421. t.Errorf("请求方法不匹配: got %s", r.Method)
  422. }
  423. // 验证 Content-Type
  424. if ct := r.Header.Get("Content-Type"); ct != "application/x-www-form-urlencoded" {
  425. t.Errorf("Content-Type 不匹配: got %s", ct)
  426. }
  427. // 验证请求体参数
  428. if err := r.ParseForm(); err != nil {
  429. t.Fatalf("解析表单失败: %v", err)
  430. }
  431. if r.FormValue("client_id") != ClientID {
  432. t.Errorf("client_id 不匹配: got %s", r.FormValue("client_id"))
  433. }
  434. if r.FormValue("client_secret") != "test-secret" {
  435. t.Errorf("client_secret 不匹配: got %s", r.FormValue("client_secret"))
  436. }
  437. if r.FormValue("code") != "auth-code" {
  438. t.Errorf("code 不匹配: got %s", r.FormValue("code"))
  439. }
  440. if r.FormValue("code_verifier") != "verifier123" {
  441. t.Errorf("code_verifier 不匹配: got %s", r.FormValue("code_verifier"))
  442. }
  443. if r.FormValue("grant_type") != "authorization_code" {
  444. t.Errorf("grant_type 不匹配: got %s", r.FormValue("grant_type"))
  445. }
  446. w.Header().Set("Content-Type", "application/json")
  447. w.WriteHeader(http.StatusOK)
  448. _ = json.NewEncoder(w).Encode(TokenResponse{
  449. AccessToken: "access-tok",
  450. ExpiresIn: 3600,
  451. TokenType: "Bearer",
  452. RefreshToken: "refresh-tok",
  453. })
  454. }))
  455. defer server.Close()
  456. // 临时替换 TokenURL(该函数直接使用常量,需要我们通过构建自定义 client 来绕过)
  457. // 由于 ExchangeCode 硬编码了 TokenURL,我们需要直接测试 HTTP client 的行为
  458. // 这里通过构造一个直接调用 mock server 的测试
  459. client := &Client{httpClient: server.Client()}
  460. // 由于 ExchangeCode 使用硬编码的 TokenURL,我们无法直接注入 mock server URL
  461. // 需要使用 httptest 的 Transport 重定向
  462. originalTokenURL := TokenURL
  463. // 我们改为直接构造请求来测试逻辑
  464. _ = originalTokenURL
  465. _ = client
  466. // 改用直接构造请求测试 mock server 响应
  467. ctx := context.Background()
  468. params := url.Values{}
  469. params.Set("client_id", ClientID)
  470. params.Set("client_secret", "test-secret")
  471. params.Set("code", "auth-code")
  472. params.Set("redirect_uri", RedirectURI)
  473. params.Set("grant_type", "authorization_code")
  474. params.Set("code_verifier", "verifier123")
  475. req, err := http.NewRequestWithContext(ctx, http.MethodPost, server.URL, strings.NewReader(params.Encode()))
  476. if err != nil {
  477. t.Fatalf("创建请求失败: %v", err)
  478. }
  479. req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
  480. resp, err := server.Client().Do(req)
  481. if err != nil {
  482. t.Fatalf("请求失败: %v", err)
  483. }
  484. defer func() { _ = resp.Body.Close() }()
  485. if resp.StatusCode != http.StatusOK {
  486. t.Fatalf("状态码不匹配: got %d", resp.StatusCode)
  487. }
  488. var tokenResp TokenResponse
  489. if err := json.NewDecoder(resp.Body).Decode(&tokenResp); err != nil {
  490. t.Fatalf("解码失败: %v", err)
  491. }
  492. if tokenResp.AccessToken != "access-tok" {
  493. t.Errorf("AccessToken 不匹配: got %s", tokenResp.AccessToken)
  494. }
  495. if tokenResp.RefreshToken != "refresh-tok" {
  496. t.Errorf("RefreshToken 不匹配: got %s", tokenResp.RefreshToken)
  497. }
  498. }
  499. func TestClient_ExchangeCode_无ClientSecret(t *testing.T) {
  500. old := defaultClientSecret
  501. defaultClientSecret = ""
  502. t.Cleanup(func() { defaultClientSecret = old })
  503. client := mustNewClient(t, "")
  504. _, err := client.ExchangeCode(context.Background(), "code", "verifier")
  505. if err == nil {
  506. t.Fatal("缺少 client_secret 时应返回错误")
  507. }
  508. if !strings.Contains(err.Error(), AntigravityOAuthClientSecretEnv) {
  509. t.Errorf("错误信息应包含环境变量名: got %s", err.Error())
  510. }
  511. }
  512. func TestClient_ExchangeCode_服务器返回错误(t *testing.T) {
  513. old := defaultClientSecret
  514. defaultClientSecret = "test-secret"
  515. t.Cleanup(func() { defaultClientSecret = old })
  516. server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
  517. w.WriteHeader(http.StatusBadRequest)
  518. _, _ = w.Write([]byte(`{"error":"invalid_grant"}`))
  519. }))
  520. defer server.Close()
  521. // 直接测试 mock server 的错误响应
  522. resp, err := server.Client().Get(server.URL)
  523. if err != nil {
  524. t.Fatalf("请求失败: %v", err)
  525. }
  526. defer func() { _ = resp.Body.Close() }()
  527. if resp.StatusCode != http.StatusBadRequest {
  528. t.Errorf("状态码不匹配: got %d, want 400", resp.StatusCode)
  529. }
  530. }
  531. // ---------------------------------------------------------------------------
  532. // Client.RefreshToken (使用 httptest)
  533. // ---------------------------------------------------------------------------
  534. func TestClient_RefreshToken_MockServer(t *testing.T) {
  535. old := defaultClientSecret
  536. defaultClientSecret = "test-secret"
  537. t.Cleanup(func() { defaultClientSecret = old })
  538. server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
  539. if r.Method != http.MethodPost {
  540. t.Errorf("请求方法不匹配: got %s", r.Method)
  541. }
  542. if err := r.ParseForm(); err != nil {
  543. t.Fatalf("解析表单失败: %v", err)
  544. }
  545. if r.FormValue("grant_type") != "refresh_token" {
  546. t.Errorf("grant_type 不匹配: got %s", r.FormValue("grant_type"))
  547. }
  548. if r.FormValue("refresh_token") != "old-refresh-tok" {
  549. t.Errorf("refresh_token 不匹配: got %s", r.FormValue("refresh_token"))
  550. }
  551. w.Header().Set("Content-Type", "application/json")
  552. w.WriteHeader(http.StatusOK)
  553. _ = json.NewEncoder(w).Encode(TokenResponse{
  554. AccessToken: "new-access-tok",
  555. ExpiresIn: 3600,
  556. TokenType: "Bearer",
  557. })
  558. }))
  559. defer server.Close()
  560. ctx := context.Background()
  561. params := url.Values{}
  562. params.Set("client_id", ClientID)
  563. params.Set("client_secret", "test-secret")
  564. params.Set("refresh_token", "old-refresh-tok")
  565. params.Set("grant_type", "refresh_token")
  566. req, err := http.NewRequestWithContext(ctx, http.MethodPost, server.URL, strings.NewReader(params.Encode()))
  567. if err != nil {
  568. t.Fatalf("创建请求失败: %v", err)
  569. }
  570. req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
  571. resp, err := server.Client().Do(req)
  572. if err != nil {
  573. t.Fatalf("请求失败: %v", err)
  574. }
  575. defer func() { _ = resp.Body.Close() }()
  576. if resp.StatusCode != http.StatusOK {
  577. t.Fatalf("状态码不匹配: got %d", resp.StatusCode)
  578. }
  579. var tokenResp TokenResponse
  580. if err := json.NewDecoder(resp.Body).Decode(&tokenResp); err != nil {
  581. t.Fatalf("解码失败: %v", err)
  582. }
  583. if tokenResp.AccessToken != "new-access-tok" {
  584. t.Errorf("AccessToken 不匹配: got %s", tokenResp.AccessToken)
  585. }
  586. }
  587. func TestClient_RefreshToken_无ClientSecret(t *testing.T) {
  588. old := defaultClientSecret
  589. defaultClientSecret = ""
  590. t.Cleanup(func() { defaultClientSecret = old })
  591. client := mustNewClient(t, "")
  592. _, err := client.RefreshToken(context.Background(), "refresh-tok")
  593. if err == nil {
  594. t.Fatal("缺少 client_secret 时应返回错误")
  595. }
  596. }
  597. // ---------------------------------------------------------------------------
  598. // Client.GetUserInfo (使用 httptest)
  599. // ---------------------------------------------------------------------------
  600. func TestClient_GetUserInfo_成功(t *testing.T) {
  601. server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
  602. if r.Method != http.MethodGet {
  603. t.Errorf("请求方法不匹配: got %s", r.Method)
  604. }
  605. auth := r.Header.Get("Authorization")
  606. if auth != "Bearer test-access-token" {
  607. t.Errorf("Authorization 不匹配: got %s", auth)
  608. }
  609. w.Header().Set("Content-Type", "application/json")
  610. w.WriteHeader(http.StatusOK)
  611. _ = json.NewEncoder(w).Encode(UserInfo{
  612. Email: "user@example.com",
  613. Name: "Test User",
  614. GivenName: "Test",
  615. FamilyName: "User",
  616. Picture: "https://example.com/photo.jpg",
  617. })
  618. }))
  619. defer server.Close()
  620. // 直接通过 mock server 测试 GetUserInfo 的行为逻辑
  621. ctx := context.Background()
  622. req, err := http.NewRequestWithContext(ctx, http.MethodGet, server.URL, nil)
  623. if err != nil {
  624. t.Fatalf("创建请求失败: %v", err)
  625. }
  626. req.Header.Set("Authorization", "Bearer test-access-token")
  627. resp, err := server.Client().Do(req)
  628. if err != nil {
  629. t.Fatalf("请求失败: %v", err)
  630. }
  631. defer func() { _ = resp.Body.Close() }()
  632. if resp.StatusCode != http.StatusOK {
  633. t.Fatalf("状态码不匹配: got %d", resp.StatusCode)
  634. }
  635. var userInfo UserInfo
  636. if err := json.NewDecoder(resp.Body).Decode(&userInfo); err != nil {
  637. t.Fatalf("解码失败: %v", err)
  638. }
  639. if userInfo.Email != "user@example.com" {
  640. t.Errorf("Email 不匹配: got %s", userInfo.Email)
  641. }
  642. if userInfo.Name != "Test User" {
  643. t.Errorf("Name 不匹配: got %s", userInfo.Name)
  644. }
  645. }
  646. func TestClient_GetUserInfo_服务器返回错误(t *testing.T) {
  647. server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
  648. w.WriteHeader(http.StatusUnauthorized)
  649. _, _ = w.Write([]byte(`{"error":"invalid_token"}`))
  650. }))
  651. defer server.Close()
  652. resp, err := server.Client().Get(server.URL)
  653. if err != nil {
  654. t.Fatalf("请求失败: %v", err)
  655. }
  656. defer func() { _ = resp.Body.Close() }()
  657. if resp.StatusCode != http.StatusUnauthorized {
  658. t.Errorf("状态码不匹配: got %d, want 401", resp.StatusCode)
  659. }
  660. }
  661. // ---------------------------------------------------------------------------
  662. // TokenResponse / UserInfo JSON 序列化
  663. // ---------------------------------------------------------------------------
  664. func TestTokenResponse_JSON序列化(t *testing.T) {
  665. jsonData := `{"access_token":"at","expires_in":3600,"token_type":"Bearer","scope":"openid","refresh_token":"rt"}`
  666. var resp TokenResponse
  667. if err := json.Unmarshal([]byte(jsonData), &resp); err != nil {
  668. t.Fatalf("反序列化失败: %v", err)
  669. }
  670. if resp.AccessToken != "at" {
  671. t.Errorf("AccessToken 不匹配: got %s", resp.AccessToken)
  672. }
  673. if resp.ExpiresIn != 3600 {
  674. t.Errorf("ExpiresIn 不匹配: got %d", resp.ExpiresIn)
  675. }
  676. if resp.RefreshToken != "rt" {
  677. t.Errorf("RefreshToken 不匹配: got %s", resp.RefreshToken)
  678. }
  679. }
  680. func TestUserInfo_JSON序列化(t *testing.T) {
  681. jsonData := `{"email":"a@b.com","name":"Alice"}`
  682. var info UserInfo
  683. if err := json.Unmarshal([]byte(jsonData), &info); err != nil {
  684. t.Fatalf("反序列化失败: %v", err)
  685. }
  686. if info.Email != "a@b.com" {
  687. t.Errorf("Email 不匹配: got %s", info.Email)
  688. }
  689. if info.Name != "Alice" {
  690. t.Errorf("Name 不匹配: got %s", info.Name)
  691. }
  692. }
  693. // ---------------------------------------------------------------------------
  694. // LoadCodeAssistResponse JSON 序列化
  695. // ---------------------------------------------------------------------------
  696. func TestLoadCodeAssistResponse_完整JSON(t *testing.T) {
  697. jsonData := `{
  698. "cloudaicompanionProject": "proj-123",
  699. "currentTier": "free-tier",
  700. "paidTier": {"id": "g1-pro-tier", "name": "Pro"},
  701. "ineligibleTiers": [{"tier": {"id": "g1-ultra-tier"}, "reasonCode": "INELIGIBLE_ACCOUNT"}]
  702. }`
  703. var resp LoadCodeAssistResponse
  704. if err := json.Unmarshal([]byte(jsonData), &resp); err != nil {
  705. t.Fatalf("反序列化失败: %v", err)
  706. }
  707. if resp.CloudAICompanionProject != "proj-123" {
  708. t.Errorf("CloudAICompanionProject 不匹配: got %s", resp.CloudAICompanionProject)
  709. }
  710. if resp.GetTier() != "g1-pro-tier" {
  711. t.Errorf("GetTier 不匹配: got %s", resp.GetTier())
  712. }
  713. if len(resp.IneligibleTiers) != 1 {
  714. t.Fatalf("IneligibleTiers 数量不匹配: got %d", len(resp.IneligibleTiers))
  715. }
  716. if resp.IneligibleTiers[0].ReasonCode != "INELIGIBLE_ACCOUNT" {
  717. t.Errorf("ReasonCode 不匹配: got %s", resp.IneligibleTiers[0].ReasonCode)
  718. }
  719. }
  720. // ===========================================================================
  721. // 以下为新增测试:真正调用 Client 方法,通过 RoundTripper 拦截 HTTP 请求
  722. // ===========================================================================
  723. // redirectRoundTripper 将请求中特定前缀的 URL 重定向到 httptest server
  724. type redirectRoundTripper struct {
  725. // 原始 URL 前缀 -> 替换目标 URL 的映射
  726. redirects map[string]string
  727. transport http.RoundTripper
  728. }
  729. type roundTripperFunc func(*http.Request) (*http.Response, error)
  730. func (f roundTripperFunc) RoundTrip(req *http.Request) (*http.Response, error) {
  731. return f(req)
  732. }
  733. func (rt *redirectRoundTripper) RoundTrip(req *http.Request) (*http.Response, error) {
  734. originalURL := req.URL.String()
  735. for prefix, target := range rt.redirects {
  736. if strings.HasPrefix(originalURL, prefix) {
  737. newURL := target + strings.TrimPrefix(originalURL, prefix)
  738. parsed, err := url.Parse(newURL)
  739. if err != nil {
  740. return nil, err
  741. }
  742. req.URL = parsed
  743. break
  744. }
  745. }
  746. if rt.transport == nil {
  747. return http.DefaultTransport.RoundTrip(req)
  748. }
  749. return rt.transport.RoundTrip(req)
  750. }
  751. // newTestClientWithRedirect 创建一个 Client,将指定 URL 前缀的请求重定向到 mock server
  752. func newTestClientWithRedirect(redirects map[string]string) *Client {
  753. return &Client{
  754. httpClient: &http.Client{
  755. Timeout: 10 * time.Second,
  756. Transport: &redirectRoundTripper{
  757. redirects: redirects,
  758. },
  759. },
  760. }
  761. }
  762. // ---------------------------------------------------------------------------
  763. // Client.ExchangeCode - 真正调用方法的测试
  764. // ---------------------------------------------------------------------------
  765. func TestClient_ExchangeCode_Success_RealCall(t *testing.T) {
  766. old := defaultClientSecret
  767. defaultClientSecret = "test-secret"
  768. t.Cleanup(func() { defaultClientSecret = old })
  769. server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
  770. if r.Method != http.MethodPost {
  771. t.Errorf("请求方法不匹配: got %s, want POST", r.Method)
  772. }
  773. if ct := r.Header.Get("Content-Type"); ct != "application/x-www-form-urlencoded" {
  774. t.Errorf("Content-Type 不匹配: got %s", ct)
  775. }
  776. if err := r.ParseForm(); err != nil {
  777. t.Fatalf("解析表单失败: %v", err)
  778. }
  779. if r.FormValue("client_id") != ClientID {
  780. t.Errorf("client_id 不匹配: got %s", r.FormValue("client_id"))
  781. }
  782. if r.FormValue("client_secret") != "test-secret" {
  783. t.Errorf("client_secret 不匹配: got %s", r.FormValue("client_secret"))
  784. }
  785. if r.FormValue("code") != "test-auth-code" {
  786. t.Errorf("code 不匹配: got %s", r.FormValue("code"))
  787. }
  788. if r.FormValue("code_verifier") != "test-verifier" {
  789. t.Errorf("code_verifier 不匹配: got %s", r.FormValue("code_verifier"))
  790. }
  791. if r.FormValue("grant_type") != "authorization_code" {
  792. t.Errorf("grant_type 不匹配: got %s", r.FormValue("grant_type"))
  793. }
  794. if r.FormValue("redirect_uri") != RedirectURI {
  795. t.Errorf("redirect_uri 不匹配: got %s", r.FormValue("redirect_uri"))
  796. }
  797. w.Header().Set("Content-Type", "application/json")
  798. w.WriteHeader(http.StatusOK)
  799. _ = json.NewEncoder(w).Encode(TokenResponse{
  800. AccessToken: "new-access-token",
  801. ExpiresIn: 3600,
  802. TokenType: "Bearer",
  803. Scope: "openid email",
  804. RefreshToken: "new-refresh-token",
  805. })
  806. }))
  807. defer server.Close()
  808. client := newTestClientWithRedirect(map[string]string{
  809. TokenURL: server.URL,
  810. })
  811. tokenResp, err := client.ExchangeCode(context.Background(), "test-auth-code", "test-verifier")
  812. if err != nil {
  813. t.Fatalf("ExchangeCode 失败: %v", err)
  814. }
  815. if tokenResp.AccessToken != "new-access-token" {
  816. t.Errorf("AccessToken 不匹配: got %s, want new-access-token", tokenResp.AccessToken)
  817. }
  818. if tokenResp.RefreshToken != "new-refresh-token" {
  819. t.Errorf("RefreshToken 不匹配: got %s, want new-refresh-token", tokenResp.RefreshToken)
  820. }
  821. if tokenResp.ExpiresIn != 3600 {
  822. t.Errorf("ExpiresIn 不匹配: got %d, want 3600", tokenResp.ExpiresIn)
  823. }
  824. if tokenResp.TokenType != "Bearer" {
  825. t.Errorf("TokenType 不匹配: got %s, want Bearer", tokenResp.TokenType)
  826. }
  827. if tokenResp.Scope != "openid email" {
  828. t.Errorf("Scope 不匹配: got %s, want openid email", tokenResp.Scope)
  829. }
  830. }
  831. func TestClient_ExchangeCode_ServerError_RealCall(t *testing.T) {
  832. old := defaultClientSecret
  833. defaultClientSecret = "test-secret"
  834. t.Cleanup(func() { defaultClientSecret = old })
  835. server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
  836. w.WriteHeader(http.StatusBadRequest)
  837. _, _ = w.Write([]byte(`{"error":"invalid_grant","error_description":"code expired"}`))
  838. }))
  839. defer server.Close()
  840. client := newTestClientWithRedirect(map[string]string{
  841. TokenURL: server.URL,
  842. })
  843. _, err := client.ExchangeCode(context.Background(), "expired-code", "verifier")
  844. if err == nil {
  845. t.Fatal("服务器返回 400 时应返回错误")
  846. }
  847. if !strings.Contains(err.Error(), "token 交换失败") {
  848. t.Errorf("错误信息应包含 'token 交换失败': got %s", err.Error())
  849. }
  850. if !strings.Contains(err.Error(), "400") {
  851. t.Errorf("错误信息应包含状态码 400: got %s", err.Error())
  852. }
  853. }
  854. func TestClient_ExchangeCode_InvalidJSON_RealCall(t *testing.T) {
  855. old := defaultClientSecret
  856. defaultClientSecret = "test-secret"
  857. t.Cleanup(func() { defaultClientSecret = old })
  858. server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
  859. w.Header().Set("Content-Type", "application/json")
  860. w.WriteHeader(http.StatusOK)
  861. _, _ = w.Write([]byte(`{invalid json`))
  862. }))
  863. defer server.Close()
  864. client := newTestClientWithRedirect(map[string]string{
  865. TokenURL: server.URL,
  866. })
  867. _, err := client.ExchangeCode(context.Background(), "code", "verifier")
  868. if err == nil {
  869. t.Fatal("无效 JSON 响应应返回错误")
  870. }
  871. if !strings.Contains(err.Error(), "token 解析失败") {
  872. t.Errorf("错误信息应包含 'token 解析失败': got %s", err.Error())
  873. }
  874. }
  875. func TestClient_ExchangeCode_ContextCanceled_RealCall(t *testing.T) {
  876. old := defaultClientSecret
  877. defaultClientSecret = "test-secret"
  878. t.Cleanup(func() { defaultClientSecret = old })
  879. server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
  880. time.Sleep(5 * time.Second) // 模拟慢响应
  881. w.WriteHeader(http.StatusOK)
  882. }))
  883. defer server.Close()
  884. client := newTestClientWithRedirect(map[string]string{
  885. TokenURL: server.URL,
  886. })
  887. ctx, cancel := context.WithCancel(context.Background())
  888. cancel() // 立即取消
  889. _, err := client.ExchangeCode(ctx, "code", "verifier")
  890. if err == nil {
  891. t.Fatal("context 取消时应返回错误")
  892. }
  893. }
  894. // ---------------------------------------------------------------------------
  895. // Client.RefreshToken - 真正调用方法的测试
  896. // ---------------------------------------------------------------------------
  897. func TestClient_RefreshToken_Success_RealCall(t *testing.T) {
  898. old := defaultClientSecret
  899. defaultClientSecret = "test-secret"
  900. t.Cleanup(func() { defaultClientSecret = old })
  901. server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
  902. if r.Method != http.MethodPost {
  903. t.Errorf("请求方法不匹配: got %s, want POST", r.Method)
  904. }
  905. if err := r.ParseForm(); err != nil {
  906. t.Fatalf("解析表单失败: %v", err)
  907. }
  908. if r.FormValue("grant_type") != "refresh_token" {
  909. t.Errorf("grant_type 不匹配: got %s", r.FormValue("grant_type"))
  910. }
  911. if r.FormValue("refresh_token") != "my-refresh-token" {
  912. t.Errorf("refresh_token 不匹配: got %s", r.FormValue("refresh_token"))
  913. }
  914. if r.FormValue("client_id") != ClientID {
  915. t.Errorf("client_id 不匹配: got %s", r.FormValue("client_id"))
  916. }
  917. if r.FormValue("client_secret") != "test-secret" {
  918. t.Errorf("client_secret 不匹配: got %s", r.FormValue("client_secret"))
  919. }
  920. w.Header().Set("Content-Type", "application/json")
  921. w.WriteHeader(http.StatusOK)
  922. _ = json.NewEncoder(w).Encode(TokenResponse{
  923. AccessToken: "refreshed-access-token",
  924. ExpiresIn: 3600,
  925. TokenType: "Bearer",
  926. })
  927. }))
  928. defer server.Close()
  929. client := newTestClientWithRedirect(map[string]string{
  930. TokenURL: server.URL,
  931. })
  932. tokenResp, err := client.RefreshToken(context.Background(), "my-refresh-token")
  933. if err != nil {
  934. t.Fatalf("RefreshToken 失败: %v", err)
  935. }
  936. if tokenResp.AccessToken != "refreshed-access-token" {
  937. t.Errorf("AccessToken 不匹配: got %s, want refreshed-access-token", tokenResp.AccessToken)
  938. }
  939. if tokenResp.ExpiresIn != 3600 {
  940. t.Errorf("ExpiresIn 不匹配: got %d, want 3600", tokenResp.ExpiresIn)
  941. }
  942. }
  943. func TestClient_RefreshToken_ServerError_RealCall(t *testing.T) {
  944. old := defaultClientSecret
  945. defaultClientSecret = "test-secret"
  946. t.Cleanup(func() { defaultClientSecret = old })
  947. server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
  948. w.WriteHeader(http.StatusUnauthorized)
  949. _, _ = w.Write([]byte(`{"error":"invalid_grant","error_description":"token revoked"}`))
  950. }))
  951. defer server.Close()
  952. client := newTestClientWithRedirect(map[string]string{
  953. TokenURL: server.URL,
  954. })
  955. _, err := client.RefreshToken(context.Background(), "revoked-token")
  956. if err == nil {
  957. t.Fatal("服务器返回 401 时应返回错误")
  958. }
  959. if !strings.Contains(err.Error(), "token 刷新失败") {
  960. t.Errorf("错误信息应包含 'token 刷新失败': got %s", err.Error())
  961. }
  962. }
  963. func TestClient_RefreshToken_InvalidJSON_RealCall(t *testing.T) {
  964. old := defaultClientSecret
  965. defaultClientSecret = "test-secret"
  966. t.Cleanup(func() { defaultClientSecret = old })
  967. server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
  968. w.Header().Set("Content-Type", "application/json")
  969. w.WriteHeader(http.StatusOK)
  970. _, _ = w.Write([]byte(`not-json`))
  971. }))
  972. defer server.Close()
  973. client := newTestClientWithRedirect(map[string]string{
  974. TokenURL: server.URL,
  975. })
  976. _, err := client.RefreshToken(context.Background(), "refresh-tok")
  977. if err == nil {
  978. t.Fatal("无效 JSON 响应应返回错误")
  979. }
  980. if !strings.Contains(err.Error(), "token 解析失败") {
  981. t.Errorf("错误信息应包含 'token 解析失败': got %s", err.Error())
  982. }
  983. }
  984. func TestClient_RefreshToken_ContextCanceled_RealCall(t *testing.T) {
  985. old := defaultClientSecret
  986. defaultClientSecret = "test-secret"
  987. t.Cleanup(func() { defaultClientSecret = old })
  988. server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
  989. time.Sleep(5 * time.Second)
  990. w.WriteHeader(http.StatusOK)
  991. }))
  992. defer server.Close()
  993. client := newTestClientWithRedirect(map[string]string{
  994. TokenURL: server.URL,
  995. })
  996. ctx, cancel := context.WithCancel(context.Background())
  997. cancel()
  998. _, err := client.RefreshToken(ctx, "refresh-tok")
  999. if err == nil {
  1000. t.Fatal("context 取消时应返回错误")
  1001. }
  1002. }
  1003. // ---------------------------------------------------------------------------
  1004. // Client.GetUserInfo - 真正调用方法的测试
  1005. // ---------------------------------------------------------------------------
  1006. func TestClient_GetUserInfo_Success_RealCall(t *testing.T) {
  1007. server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
  1008. if r.Method != http.MethodGet {
  1009. t.Errorf("请求方法不匹配: got %s, want GET", r.Method)
  1010. }
  1011. auth := r.Header.Get("Authorization")
  1012. if auth != "Bearer user-access-token" {
  1013. t.Errorf("Authorization 不匹配: got %s", auth)
  1014. }
  1015. w.Header().Set("Content-Type", "application/json")
  1016. w.WriteHeader(http.StatusOK)
  1017. _ = json.NewEncoder(w).Encode(UserInfo{
  1018. Email: "test@example.com",
  1019. Name: "Test User",
  1020. GivenName: "Test",
  1021. FamilyName: "User",
  1022. Picture: "https://example.com/avatar.jpg",
  1023. })
  1024. }))
  1025. defer server.Close()
  1026. client := newTestClientWithRedirect(map[string]string{
  1027. UserInfoURL: server.URL,
  1028. })
  1029. userInfo, err := client.GetUserInfo(context.Background(), "user-access-token")
  1030. if err != nil {
  1031. t.Fatalf("GetUserInfo 失败: %v", err)
  1032. }
  1033. if userInfo.Email != "test@example.com" {
  1034. t.Errorf("Email 不匹配: got %s, want test@example.com", userInfo.Email)
  1035. }
  1036. if userInfo.Name != "Test User" {
  1037. t.Errorf("Name 不匹配: got %s, want Test User", userInfo.Name)
  1038. }
  1039. if userInfo.GivenName != "Test" {
  1040. t.Errorf("GivenName 不匹配: got %s, want Test", userInfo.GivenName)
  1041. }
  1042. if userInfo.FamilyName != "User" {
  1043. t.Errorf("FamilyName 不匹配: got %s, want User", userInfo.FamilyName)
  1044. }
  1045. if userInfo.Picture != "https://example.com/avatar.jpg" {
  1046. t.Errorf("Picture 不匹配: got %s", userInfo.Picture)
  1047. }
  1048. }
  1049. func TestClient_GetUserInfo_Unauthorized_RealCall(t *testing.T) {
  1050. server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
  1051. w.WriteHeader(http.StatusUnauthorized)
  1052. _, _ = w.Write([]byte(`{"error":"invalid_token"}`))
  1053. }))
  1054. defer server.Close()
  1055. client := newTestClientWithRedirect(map[string]string{
  1056. UserInfoURL: server.URL,
  1057. })
  1058. _, err := client.GetUserInfo(context.Background(), "bad-token")
  1059. if err == nil {
  1060. t.Fatal("服务器返回 401 时应返回错误")
  1061. }
  1062. if !strings.Contains(err.Error(), "获取用户信息失败") {
  1063. t.Errorf("错误信息应包含 '获取用户信息失败': got %s", err.Error())
  1064. }
  1065. if !strings.Contains(err.Error(), "401") {
  1066. t.Errorf("错误信息应包含状态码 401: got %s", err.Error())
  1067. }
  1068. }
  1069. func TestClient_GetUserInfo_InvalidJSON_RealCall(t *testing.T) {
  1070. server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
  1071. w.Header().Set("Content-Type", "application/json")
  1072. w.WriteHeader(http.StatusOK)
  1073. _, _ = w.Write([]byte(`{broken`))
  1074. }))
  1075. defer server.Close()
  1076. client := newTestClientWithRedirect(map[string]string{
  1077. UserInfoURL: server.URL,
  1078. })
  1079. _, err := client.GetUserInfo(context.Background(), "token")
  1080. if err == nil {
  1081. t.Fatal("无效 JSON 响应应返回错误")
  1082. }
  1083. if !strings.Contains(err.Error(), "用户信息解析失败") {
  1084. t.Errorf("错误信息应包含 '用户信息解析失败': got %s", err.Error())
  1085. }
  1086. }
  1087. func TestClient_GetUserInfo_ContextCanceled_RealCall(t *testing.T) {
  1088. server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
  1089. time.Sleep(5 * time.Second)
  1090. w.WriteHeader(http.StatusOK)
  1091. }))
  1092. defer server.Close()
  1093. client := newTestClientWithRedirect(map[string]string{
  1094. UserInfoURL: server.URL,
  1095. })
  1096. ctx, cancel := context.WithCancel(context.Background())
  1097. cancel()
  1098. _, err := client.GetUserInfo(ctx, "token")
  1099. if err == nil {
  1100. t.Fatal("context 取消时应返回错误")
  1101. }
  1102. }
  1103. // ---------------------------------------------------------------------------
  1104. // Client.LoadCodeAssist - 真正调用方法的测试
  1105. // ---------------------------------------------------------------------------
  1106. // withMockBaseURLs 临时替换 BaseURLs,测试结束后恢复
  1107. func withMockBaseURLs(t *testing.T, urls []string) {
  1108. t.Helper()
  1109. origBaseURLs := BaseURLs
  1110. origBaseURL := BaseURL
  1111. BaseURLs = urls
  1112. if len(urls) > 0 {
  1113. BaseURL = urls[0]
  1114. }
  1115. t.Cleanup(func() {
  1116. BaseURLs = origBaseURLs
  1117. BaseURL = origBaseURL
  1118. })
  1119. }
  1120. func TestClient_LoadCodeAssist_Success_RealCall(t *testing.T) {
  1121. server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
  1122. if r.Method != http.MethodPost {
  1123. t.Errorf("请求方法不匹配: got %s, want POST", r.Method)
  1124. }
  1125. if !strings.HasSuffix(r.URL.Path, "/v1internal:loadCodeAssist") {
  1126. t.Errorf("URL 路径不匹配: got %s", r.URL.Path)
  1127. }
  1128. auth := r.Header.Get("Authorization")
  1129. if auth != "Bearer test-token" {
  1130. t.Errorf("Authorization 不匹配: got %s", auth)
  1131. }
  1132. if ct := r.Header.Get("Content-Type"); ct != "application/json" {
  1133. t.Errorf("Content-Type 不匹配: got %s", ct)
  1134. }
  1135. if ua := r.Header.Get("User-Agent"); ua != GetUserAgent() {
  1136. t.Errorf("User-Agent 不匹配: got %s", ua)
  1137. }
  1138. // 验证请求体
  1139. var reqBody LoadCodeAssistRequest
  1140. if err := json.NewDecoder(r.Body).Decode(&reqBody); err != nil {
  1141. t.Fatalf("解析请求体失败: %v", err)
  1142. }
  1143. if reqBody.Metadata.IDEType != "ANTIGRAVITY" {
  1144. t.Errorf("IDEType 不匹配: got %s, want ANTIGRAVITY", reqBody.Metadata.IDEType)
  1145. }
  1146. if strings.TrimSpace(reqBody.Metadata.IDEVersion) == "" {
  1147. t.Errorf("IDEVersion 不应为空")
  1148. }
  1149. if reqBody.Metadata.IDEName != "antigravity" {
  1150. t.Errorf("IDEName 不匹配: got %s, want antigravity", reqBody.Metadata.IDEName)
  1151. }
  1152. w.Header().Set("Content-Type", "application/json")
  1153. w.WriteHeader(http.StatusOK)
  1154. _, _ = w.Write([]byte(`{
  1155. "cloudaicompanionProject": "test-project-123",
  1156. "currentTier": {"id": "free-tier", "name": "Free"},
  1157. "paidTier": {"id": "g1-pro-tier", "name": "Pro", "description": "Pro plan"}
  1158. }`))
  1159. }))
  1160. defer server.Close()
  1161. withMockBaseURLs(t, []string{server.URL})
  1162. client := mustNewClient(t, "")
  1163. resp, rawResp, err := client.LoadCodeAssist(context.Background(), "test-token")
  1164. if err != nil {
  1165. t.Fatalf("LoadCodeAssist 失败: %v", err)
  1166. }
  1167. if resp.CloudAICompanionProject != "test-project-123" {
  1168. t.Errorf("CloudAICompanionProject 不匹配: got %s", resp.CloudAICompanionProject)
  1169. }
  1170. if resp.GetTier() != "g1-pro-tier" {
  1171. t.Errorf("GetTier 不匹配: got %s, want g1-pro-tier", resp.GetTier())
  1172. }
  1173. if resp.CurrentTier == nil || resp.CurrentTier.ID != "free-tier" {
  1174. t.Errorf("CurrentTier 不匹配: got %+v", resp.CurrentTier)
  1175. }
  1176. if resp.PaidTier == nil || resp.PaidTier.ID != "g1-pro-tier" {
  1177. t.Errorf("PaidTier 不匹配: got %+v", resp.PaidTier)
  1178. }
  1179. // 验证原始 JSON map
  1180. if rawResp == nil {
  1181. t.Fatal("rawResp 不应为 nil")
  1182. }
  1183. if rawResp["cloudaicompanionProject"] != "test-project-123" {
  1184. t.Errorf("rawResp cloudaicompanionProject 不匹配: got %v", rawResp["cloudaicompanionProject"])
  1185. }
  1186. }
  1187. func TestClient_LoadCodeAssist_HTTPError_RealCall(t *testing.T) {
  1188. server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
  1189. w.WriteHeader(http.StatusForbidden)
  1190. _, _ = w.Write([]byte(`{"error":"forbidden"}`))
  1191. }))
  1192. defer server.Close()
  1193. withMockBaseURLs(t, []string{server.URL})
  1194. client := mustNewClient(t, "")
  1195. _, _, err := client.LoadCodeAssist(context.Background(), "bad-token")
  1196. if err == nil {
  1197. t.Fatal("服务器返回 403 时应返回错误")
  1198. }
  1199. if !strings.Contains(err.Error(), "loadCodeAssist 失败") {
  1200. t.Errorf("错误信息应包含 'loadCodeAssist 失败': got %s", err.Error())
  1201. }
  1202. if !strings.Contains(err.Error(), "403") {
  1203. t.Errorf("错误信息应包含状态码 403: got %s", err.Error())
  1204. }
  1205. }
  1206. func TestClient_LoadCodeAssist_InvalidJSON_RealCall(t *testing.T) {
  1207. server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
  1208. w.Header().Set("Content-Type", "application/json")
  1209. w.WriteHeader(http.StatusOK)
  1210. _, _ = w.Write([]byte(`{not valid json!!!`))
  1211. }))
  1212. defer server.Close()
  1213. withMockBaseURLs(t, []string{server.URL})
  1214. client := mustNewClient(t, "")
  1215. _, _, err := client.LoadCodeAssist(context.Background(), "token")
  1216. if err == nil {
  1217. t.Fatal("无效 JSON 响应应返回错误")
  1218. }
  1219. if !strings.Contains(err.Error(), "响应解析失败") {
  1220. t.Errorf("错误信息应包含 '响应解析失败': got %s", err.Error())
  1221. }
  1222. }
  1223. func TestClient_LoadCodeAssist_URLFallback_RealCall(t *testing.T) {
  1224. // 第一个 server 返回 500,第二个 server 返回成功
  1225. callCount := 0
  1226. server1 := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
  1227. callCount++
  1228. w.WriteHeader(http.StatusInternalServerError)
  1229. _, _ = w.Write([]byte(`{"error":"internal"}`))
  1230. }))
  1231. defer server1.Close()
  1232. server2 := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
  1233. callCount++
  1234. w.Header().Set("Content-Type", "application/json")
  1235. w.WriteHeader(http.StatusOK)
  1236. _, _ = w.Write([]byte(`{
  1237. "cloudaicompanionProject": "fallback-project",
  1238. "currentTier": {"id": "free-tier", "name": "Free"}
  1239. }`))
  1240. }))
  1241. defer server2.Close()
  1242. withMockBaseURLs(t, []string{server1.URL, server2.URL})
  1243. client := mustNewClient(t, "")
  1244. resp, _, err := client.LoadCodeAssist(context.Background(), "token")
  1245. if err != nil {
  1246. t.Fatalf("LoadCodeAssist 应在 fallback 后成功: %v", err)
  1247. }
  1248. if resp.CloudAICompanionProject != "fallback-project" {
  1249. t.Errorf("CloudAICompanionProject 不匹配: got %s", resp.CloudAICompanionProject)
  1250. }
  1251. if callCount != 2 {
  1252. t.Errorf("应该调用了 2 个 server,实际调用 %d 次", callCount)
  1253. }
  1254. }
  1255. func TestClient_LoadCodeAssist_AllURLsFail_RealCall(t *testing.T) {
  1256. server1 := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
  1257. w.WriteHeader(http.StatusServiceUnavailable)
  1258. _, _ = w.Write([]byte(`{"error":"unavailable"}`))
  1259. }))
  1260. defer server1.Close()
  1261. server2 := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
  1262. w.WriteHeader(http.StatusBadGateway)
  1263. _, _ = w.Write([]byte(`{"error":"bad_gateway"}`))
  1264. }))
  1265. defer server2.Close()
  1266. withMockBaseURLs(t, []string{server1.URL, server2.URL})
  1267. client := mustNewClient(t, "")
  1268. _, _, err := client.LoadCodeAssist(context.Background(), "token")
  1269. if err == nil {
  1270. t.Fatal("所有 URL 都失败时应返回错误")
  1271. }
  1272. }
  1273. func TestClient_LoadCodeAssist_ContextCanceled_RealCall(t *testing.T) {
  1274. server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
  1275. time.Sleep(5 * time.Second)
  1276. w.WriteHeader(http.StatusOK)
  1277. }))
  1278. defer server.Close()
  1279. withMockBaseURLs(t, []string{server.URL})
  1280. client := mustNewClient(t, "")
  1281. ctx, cancel := context.WithCancel(context.Background())
  1282. cancel()
  1283. _, _, err := client.LoadCodeAssist(ctx, "token")
  1284. if err == nil {
  1285. t.Fatal("context 取消时应返回错误")
  1286. }
  1287. }
  1288. // ---------------------------------------------------------------------------
  1289. // Client.FetchAvailableModels - 真正调用方法的测试
  1290. // ---------------------------------------------------------------------------
  1291. func TestClient_FetchAvailableModels_Success_RealCall(t *testing.T) {
  1292. server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
  1293. if r.Method != http.MethodPost {
  1294. t.Errorf("请求方法不匹配: got %s, want POST", r.Method)
  1295. }
  1296. if !strings.HasSuffix(r.URL.Path, "/v1internal:fetchAvailableModels") {
  1297. t.Errorf("URL 路径不匹配: got %s", r.URL.Path)
  1298. }
  1299. auth := r.Header.Get("Authorization")
  1300. if auth != "Bearer test-token" {
  1301. t.Errorf("Authorization 不匹配: got %s", auth)
  1302. }
  1303. if ct := r.Header.Get("Content-Type"); ct != "application/json" {
  1304. t.Errorf("Content-Type 不匹配: got %s", ct)
  1305. }
  1306. if ua := r.Header.Get("User-Agent"); ua != GetUserAgent() {
  1307. t.Errorf("User-Agent 不匹配: got %s", ua)
  1308. }
  1309. // 验证请求体
  1310. var reqBody FetchAvailableModelsRequest
  1311. if err := json.NewDecoder(r.Body).Decode(&reqBody); err != nil {
  1312. t.Fatalf("解析请求体失败: %v", err)
  1313. }
  1314. if reqBody.Project != "project-abc" {
  1315. t.Errorf("Project 不匹配: got %s, want project-abc", reqBody.Project)
  1316. }
  1317. w.Header().Set("Content-Type", "application/json")
  1318. w.WriteHeader(http.StatusOK)
  1319. _, _ = w.Write([]byte(`{
  1320. "models": {
  1321. "gemini-2.0-flash": {
  1322. "quotaInfo": {
  1323. "remainingFraction": 0.85,
  1324. "resetTime": "2025-01-01T00:00:00Z"
  1325. }
  1326. },
  1327. "gemini-2.5-pro": {
  1328. "quotaInfo": {
  1329. "remainingFraction": 0.5
  1330. }
  1331. }
  1332. }
  1333. }`))
  1334. }))
  1335. defer server.Close()
  1336. withMockBaseURLs(t, []string{server.URL})
  1337. client := mustNewClient(t, "")
  1338. resp, rawResp, err := client.FetchAvailableModels(context.Background(), "test-token", "project-abc", defaultFetchAvailableModelsBodyLimit)
  1339. if err != nil {
  1340. t.Fatalf("FetchAvailableModels 失败: %v", err)
  1341. }
  1342. if resp.Models == nil {
  1343. t.Fatal("Models 不应为 nil")
  1344. }
  1345. if len(resp.Models) != 2 {
  1346. t.Errorf("Models 数量不匹配: got %d, want 2", len(resp.Models))
  1347. }
  1348. flashModel, ok := resp.Models["gemini-2.0-flash"]
  1349. if !ok {
  1350. t.Fatal("缺少 gemini-2.0-flash 模型")
  1351. }
  1352. if flashModel.QuotaInfo == nil {
  1353. t.Fatal("gemini-2.0-flash QuotaInfo 不应为 nil")
  1354. }
  1355. if flashModel.QuotaInfo.RemainingFraction != 0.85 {
  1356. t.Errorf("RemainingFraction 不匹配: got %f, want 0.85", flashModel.QuotaInfo.RemainingFraction)
  1357. }
  1358. if flashModel.QuotaInfo.ResetTime != "2025-01-01T00:00:00Z" {
  1359. t.Errorf("ResetTime 不匹配: got %s", flashModel.QuotaInfo.ResetTime)
  1360. }
  1361. proModel, ok := resp.Models["gemini-2.5-pro"]
  1362. if !ok {
  1363. t.Fatal("缺少 gemini-2.5-pro 模型")
  1364. }
  1365. if proModel.QuotaInfo == nil {
  1366. t.Fatal("gemini-2.5-pro QuotaInfo 不应为 nil")
  1367. }
  1368. if proModel.QuotaInfo.RemainingFraction != 0.5 {
  1369. t.Errorf("RemainingFraction 不匹配: got %f, want 0.5", proModel.QuotaInfo.RemainingFraction)
  1370. }
  1371. // 验证原始 JSON map
  1372. if rawResp == nil {
  1373. t.Fatal("rawResp 不应为 nil")
  1374. }
  1375. if rawResp["models"] == nil {
  1376. t.Error("rawResp models 不应为 nil")
  1377. }
  1378. }
  1379. func TestClient_FetchAvailableModels_RejectsNonPositiveBodyLimit(t *testing.T) {
  1380. client := mustNewClient(t, "")
  1381. _, _, err := client.FetchAvailableModels(context.Background(), "token", "proj", 0)
  1382. if err == nil || !strings.Contains(err.Error(), "body limit must be positive") {
  1383. t.Fatalf("非正读取上限应返回调用错误: %v", err)
  1384. }
  1385. }
  1386. func TestClient_FetchAvailableModels_UsesConfiguredBodyLimit(t *testing.T) {
  1387. server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
  1388. w.Header().Set("Content-Type", "application/json")
  1389. w.WriteHeader(http.StatusOK)
  1390. _, _ = w.Write([]byte(`{"models":{"model-a":{}}}`))
  1391. }))
  1392. defer server.Close()
  1393. withMockBaseURLs(t, []string{server.URL})
  1394. client := mustNewClient(t, "")
  1395. _, _, err := client.FetchAvailableModels(context.Background(), "token", "proj", 8)
  1396. if err == nil {
  1397. t.Fatal("响应超过配置上限时应返回错误")
  1398. }
  1399. if !strings.Contains(err.Error(), "响应超过 8 字节") {
  1400. t.Fatalf("错误应包含配置上限: %v", err)
  1401. }
  1402. }
  1403. func TestClient_FetchAvailableModels_HTTPError_RealCall(t *testing.T) {
  1404. server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
  1405. w.WriteHeader(http.StatusForbidden)
  1406. _, _ = w.Write([]byte(`{"error":"forbidden"}`))
  1407. }))
  1408. defer server.Close()
  1409. withMockBaseURLs(t, []string{server.URL})
  1410. client := mustNewClient(t, "")
  1411. _, _, err := client.FetchAvailableModels(context.Background(), "bad-token", "proj", defaultFetchAvailableModelsBodyLimit)
  1412. if err == nil {
  1413. t.Fatal("服务器返回 403 时应返回错误")
  1414. }
  1415. if !strings.Contains(err.Error(), "fetchAvailableModels 失败") {
  1416. t.Errorf("错误信息应包含 'fetchAvailableModels 失败': got %s", err.Error())
  1417. }
  1418. }
  1419. func TestClient_FetchAvailableModels_InvalidJSON_RealCall(t *testing.T) {
  1420. server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
  1421. w.Header().Set("Content-Type", "application/json")
  1422. w.WriteHeader(http.StatusOK)
  1423. _, _ = w.Write([]byte(`<<<not json>>>`))
  1424. }))
  1425. defer server.Close()
  1426. withMockBaseURLs(t, []string{server.URL})
  1427. client := mustNewClient(t, "")
  1428. _, _, err := client.FetchAvailableModels(context.Background(), "token", "proj", defaultFetchAvailableModelsBodyLimit)
  1429. if err == nil {
  1430. t.Fatal("无效 JSON 响应应返回错误")
  1431. }
  1432. if !strings.Contains(err.Error(), "响应解析失败") {
  1433. t.Errorf("错误信息应包含 '响应解析失败': got %s", err.Error())
  1434. }
  1435. }
  1436. func TestClient_FetchAvailableModels_URLFallback_RealCall(t *testing.T) {
  1437. callCount := 0
  1438. // 第一个 server 返回 429,第二个 server 返回成功
  1439. server1 := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
  1440. callCount++
  1441. w.WriteHeader(http.StatusTooManyRequests)
  1442. _, _ = w.Write([]byte(`{"error":"rate_limited"}`))
  1443. }))
  1444. defer server1.Close()
  1445. server2 := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
  1446. callCount++
  1447. w.Header().Set("Content-Type", "application/json")
  1448. w.WriteHeader(http.StatusOK)
  1449. _, _ = w.Write([]byte(`{"models": {"model-a": {}}}`))
  1450. }))
  1451. defer server2.Close()
  1452. withMockBaseURLs(t, []string{server1.URL, server2.URL})
  1453. client := mustNewClient(t, "")
  1454. resp, _, err := client.FetchAvailableModels(context.Background(), "token", "proj", defaultFetchAvailableModelsBodyLimit)
  1455. if err != nil {
  1456. t.Fatalf("FetchAvailableModels 应在 fallback 后成功: %v", err)
  1457. }
  1458. if _, ok := resp.Models["model-a"]; !ok {
  1459. t.Error("应返回 fallback server 的模型")
  1460. }
  1461. if callCount != 2 {
  1462. t.Errorf("应该调用了 2 个 server,实际调用 %d 次", callCount)
  1463. }
  1464. }
  1465. func TestClient_FetchAvailableModels_AllURLsFail_RealCall(t *testing.T) {
  1466. server1 := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
  1467. w.WriteHeader(http.StatusNotFound)
  1468. _, _ = w.Write([]byte(`not found`))
  1469. }))
  1470. defer server1.Close()
  1471. server2 := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
  1472. w.WriteHeader(http.StatusInternalServerError)
  1473. _, _ = w.Write([]byte(`internal error`))
  1474. }))
  1475. defer server2.Close()
  1476. withMockBaseURLs(t, []string{server1.URL, server2.URL})
  1477. client := mustNewClient(t, "")
  1478. _, _, err := client.FetchAvailableModels(context.Background(), "token", "proj", defaultFetchAvailableModelsBodyLimit)
  1479. if err == nil {
  1480. t.Fatal("所有 URL 都失败时应返回错误")
  1481. }
  1482. }
  1483. func TestClient_FetchAvailableModels_ContextCanceled_RealCall(t *testing.T) {
  1484. server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
  1485. time.Sleep(5 * time.Second)
  1486. w.WriteHeader(http.StatusOK)
  1487. }))
  1488. defer server.Close()
  1489. withMockBaseURLs(t, []string{server.URL})
  1490. client := mustNewClient(t, "")
  1491. ctx, cancel := context.WithCancel(context.Background())
  1492. cancel()
  1493. _, _, err := client.FetchAvailableModels(ctx, "token", "proj", defaultFetchAvailableModelsBodyLimit)
  1494. if err == nil {
  1495. t.Fatal("context 取消时应返回错误")
  1496. }
  1497. }
  1498. func TestClient_FetchAvailableModels_EmptyModels_RealCall(t *testing.T) {
  1499. server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
  1500. w.Header().Set("Content-Type", "application/json")
  1501. w.WriteHeader(http.StatusOK)
  1502. _, _ = w.Write([]byte(`{"models": {}}`))
  1503. }))
  1504. defer server.Close()
  1505. withMockBaseURLs(t, []string{server.URL})
  1506. client := mustNewClient(t, "")
  1507. resp, rawResp, err := client.FetchAvailableModels(context.Background(), "token", "proj", defaultFetchAvailableModelsBodyLimit)
  1508. if err != nil {
  1509. t.Fatalf("FetchAvailableModels 失败: %v", err)
  1510. }
  1511. if resp.Models == nil {
  1512. t.Fatal("Models 不应为 nil")
  1513. }
  1514. if len(resp.Models) != 0 {
  1515. t.Errorf("Models 应为空: got %d", len(resp.Models))
  1516. }
  1517. if rawResp == nil {
  1518. t.Fatal("rawResp 不应为 nil")
  1519. }
  1520. }
  1521. // ---------------------------------------------------------------------------
  1522. // LoadCodeAssist 和 FetchAvailableModels 的 408 fallback 测试
  1523. // ---------------------------------------------------------------------------
  1524. func TestClient_LoadCodeAssist_408Fallback_RealCall(t *testing.T) {
  1525. server1 := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
  1526. w.WriteHeader(http.StatusRequestTimeout)
  1527. _, _ = w.Write([]byte(`timeout`))
  1528. }))
  1529. defer server1.Close()
  1530. server2 := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
  1531. w.Header().Set("Content-Type", "application/json")
  1532. w.WriteHeader(http.StatusOK)
  1533. _, _ = w.Write([]byte(`{"cloudaicompanionProject":"p2","currentTier":"free-tier"}`))
  1534. }))
  1535. defer server2.Close()
  1536. withMockBaseURLs(t, []string{server1.URL, server2.URL})
  1537. client := mustNewClient(t, "")
  1538. resp, _, err := client.LoadCodeAssist(context.Background(), "token")
  1539. if err != nil {
  1540. t.Fatalf("LoadCodeAssist 应在 408 fallback 后成功: %v", err)
  1541. }
  1542. if resp.CloudAICompanionProject != "p2" {
  1543. t.Errorf("CloudAICompanionProject 不匹配: got %s", resp.CloudAICompanionProject)
  1544. }
  1545. }
  1546. func TestClient_FetchAvailableModels_404Fallback_RealCall(t *testing.T) {
  1547. server1 := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
  1548. w.WriteHeader(http.StatusNotFound)
  1549. _, _ = w.Write([]byte(`not found`))
  1550. }))
  1551. defer server1.Close()
  1552. server2 := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
  1553. w.Header().Set("Content-Type", "application/json")
  1554. w.WriteHeader(http.StatusOK)
  1555. _, _ = w.Write([]byte(`{"models":{"m1":{"quotaInfo":{"remainingFraction":1.0}}}}`))
  1556. }))
  1557. defer server2.Close()
  1558. withMockBaseURLs(t, []string{server1.URL, server2.URL})
  1559. client := mustNewClient(t, "")
  1560. resp, _, err := client.FetchAvailableModels(context.Background(), "token", "proj", defaultFetchAvailableModelsBodyLimit)
  1561. if err != nil {
  1562. t.Fatalf("FetchAvailableModels 应在 404 fallback 后成功: %v", err)
  1563. }
  1564. if _, ok := resp.Models["m1"]; !ok {
  1565. t.Error("应返回 fallback server 的模型 m1")
  1566. }
  1567. }
  1568. func TestExtractProjectIDFromOnboardResponse(t *testing.T) {
  1569. t.Parallel()
  1570. tests := []struct {
  1571. name string
  1572. resp map[string]any
  1573. want string
  1574. }{
  1575. {
  1576. name: "nil response",
  1577. resp: nil,
  1578. want: "",
  1579. },
  1580. {
  1581. name: "empty response",
  1582. resp: map[string]any{},
  1583. want: "",
  1584. },
  1585. {
  1586. name: "project as string",
  1587. resp: map[string]any{
  1588. "cloudaicompanionProject": "my-project-123",
  1589. },
  1590. want: "my-project-123",
  1591. },
  1592. {
  1593. name: "project as string with spaces",
  1594. resp: map[string]any{
  1595. "cloudaicompanionProject": " my-project-123 ",
  1596. },
  1597. want: "my-project-123",
  1598. },
  1599. {
  1600. name: "project as map with id",
  1601. resp: map[string]any{
  1602. "cloudaicompanionProject": map[string]any{
  1603. "id": "proj-from-map",
  1604. },
  1605. },
  1606. want: "proj-from-map",
  1607. },
  1608. {
  1609. name: "project as map without id",
  1610. resp: map[string]any{
  1611. "cloudaicompanionProject": map[string]any{
  1612. "name": "some-name",
  1613. },
  1614. },
  1615. want: "",
  1616. },
  1617. {
  1618. name: "missing cloudaicompanionProject key",
  1619. resp: map[string]any{
  1620. "otherField": "value",
  1621. },
  1622. want: "",
  1623. },
  1624. }
  1625. for _, tc := range tests {
  1626. t.Run(tc.name, func(t *testing.T) {
  1627. t.Parallel()
  1628. got := extractProjectIDFromOnboardResponse(tc.resp)
  1629. if got != tc.want {
  1630. t.Fatalf("extractProjectIDFromOnboardResponse() = %q, want %q", got, tc.want)
  1631. }
  1632. })
  1633. }
  1634. }