gemini_v1beta_handler_test.go 6.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214
  1. //go:build unit
  2. package handler
  3. import (
  4. "encoding/json"
  5. "net/http"
  6. "net/http/httptest"
  7. "testing"
  8. "github.com/Wei-Shaw/sub2api/internal/pkg/antigravity"
  9. "github.com/Wei-Shaw/sub2api/internal/server/middleware"
  10. "github.com/Wei-Shaw/sub2api/internal/service"
  11. "github.com/gin-gonic/gin"
  12. "github.com/stretchr/testify/require"
  13. )
  14. func TestGeminiV1BetaListModels_AllowlistFiltersNativeResponse(t *testing.T) {
  15. body := []byte(`{"models":[{"name":"models/gemini-2.5-pro"},{"name":"models/gemini-2.5-flash"}],"nextPageToken":"next"}`)
  16. filtered, dropped, ok := filterUpstreamGeminiModelsBody(body, service.GroupModelAllowlist{Enabled: true, Models: []string{"gemini-2.5-pro"}})
  17. require.True(t, ok)
  18. require.True(t, dropped)
  19. require.JSONEq(t, `{"models":[{"name":"models/gemini-2.5-pro"}],"nextPageToken":"next"}`, string(filtered))
  20. }
  21. func TestGeminiV1BetaListModels_ForcedAntigravityAppliesAllowlist(t *testing.T) {
  22. gin.SetMode(gin.TestMode)
  23. rec := httptest.NewRecorder()
  24. c, _ := gin.CreateTestContext(rec)
  25. c.Request = httptest.NewRequest(http.MethodGet, "/antigravity/v1beta/models", nil)
  26. c.Set(string(middleware.ContextKeyAPIKey), &service.APIKey{
  27. Group: &service.Group{
  28. Platform: service.PlatformGemini,
  29. ModelAllowlist: service.GroupModelAllowlist{
  30. Enabled: true,
  31. Models: []string{"gemini-custom"},
  32. },
  33. },
  34. })
  35. c.Set(string(middleware.ContextKeyForcePlatform), service.PlatformAntigravity)
  36. (&GatewayHandler{}).GeminiV1BetaListModels(c)
  37. require.Equal(t, http.StatusOK, rec.Code)
  38. var got antigravity.GeminiModelsListResponse
  39. require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &got))
  40. require.Empty(t, got.Models)
  41. }
  42. func TestGeminiModelAllowlist_DisabledPreservesNativeResponse(t *testing.T) {
  43. body := []byte(`{"models":[{"name":"models/gemini-2.5-pro"}]}`)
  44. filtered, dropped, ok := filterUpstreamGeminiModelsBody(body, service.GroupModelAllowlist{Enabled: false, Models: []string{"other"}})
  45. require.True(t, ok)
  46. require.False(t, dropped)
  47. require.Equal(t, body, filtered)
  48. }
  49. // TestGeminiV1BetaHandler_PlatformRoutingInvariant 文档化并验证 Handler 层的平台路由逻辑不变量
  50. // 该测试确保 gemini 和 antigravity 平台的路由逻辑符合预期
  51. func TestGeminiV1BetaHandler_PlatformRoutingInvariant(t *testing.T) {
  52. tests := []struct {
  53. name string
  54. platform string
  55. expectedService string
  56. description string
  57. }{
  58. {
  59. name: "Gemini平台使用ForwardNative",
  60. platform: service.PlatformGemini,
  61. expectedService: "GeminiMessagesCompatService.ForwardNative",
  62. description: "Gemini OAuth 账户直接调用 Google API",
  63. },
  64. {
  65. name: "Antigravity平台使用ForwardGemini",
  66. platform: service.PlatformAntigravity,
  67. expectedService: "AntigravityGatewayService.ForwardGemini",
  68. description: "Antigravity 账户通过 CRS 中转,支持 Gemini 协议",
  69. },
  70. }
  71. for _, tt := range tests {
  72. t.Run(tt.name, func(t *testing.T) {
  73. // 模拟 GeminiV1BetaModels 中的路由决策 (lines 199-205 in gemini_v1beta_handler.go)
  74. var routedService string
  75. if tt.platform == service.PlatformAntigravity {
  76. routedService = "AntigravityGatewayService.ForwardGemini"
  77. } else {
  78. routedService = "GeminiMessagesCompatService.ForwardNative"
  79. }
  80. require.Equal(t, tt.expectedService, routedService,
  81. "平台 %s 应该路由到 %s: %s",
  82. tt.platform, tt.expectedService, tt.description)
  83. })
  84. }
  85. }
  86. // TestGeminiV1BetaHandler_ListModelsAntigravityFallback 验证 ListModels 的 antigravity 降级逻辑
  87. // 当没有 gemini 账户但有 antigravity 账户时,应返回静态模型列表
  88. func TestGeminiV1BetaHandler_ListModelsAntigravityFallback(t *testing.T) {
  89. tests := []struct {
  90. name string
  91. hasGeminiAccount bool
  92. hasAntigravity bool
  93. expectedBehavior string
  94. }{
  95. {
  96. name: "有Gemini账户-调用ForwardAIStudioGET",
  97. hasGeminiAccount: true,
  98. hasAntigravity: false,
  99. expectedBehavior: "forward_to_upstream",
  100. },
  101. {
  102. name: "无Gemini有Antigravity-返回静态列表",
  103. hasGeminiAccount: false,
  104. hasAntigravity: true,
  105. expectedBehavior: "static_fallback",
  106. },
  107. {
  108. name: "无任何账户-返回503",
  109. hasGeminiAccount: false,
  110. hasAntigravity: false,
  111. expectedBehavior: "service_unavailable",
  112. },
  113. }
  114. for _, tt := range tests {
  115. t.Run(tt.name, func(t *testing.T) {
  116. // 模拟 GeminiV1BetaListModels 的逻辑 (lines 33-44 in gemini_v1beta_handler.go)
  117. var behavior string
  118. if tt.hasGeminiAccount {
  119. behavior = "forward_to_upstream"
  120. } else if tt.hasAntigravity {
  121. behavior = "static_fallback"
  122. } else {
  123. behavior = "service_unavailable"
  124. }
  125. require.Equal(t, tt.expectedBehavior, behavior)
  126. })
  127. }
  128. }
  129. // TestGeminiV1BetaHandler_GetModelAntigravityFallback 验证 GetModel 的 antigravity 降级逻辑
  130. func TestGeminiV1BetaHandler_GetModelAntigravityFallback(t *testing.T) {
  131. tests := []struct {
  132. name string
  133. hasGeminiAccount bool
  134. hasAntigravity bool
  135. expectedBehavior string
  136. }{
  137. {
  138. name: "有Gemini账户-调用ForwardAIStudioGET",
  139. hasGeminiAccount: true,
  140. hasAntigravity: false,
  141. expectedBehavior: "forward_to_upstream",
  142. },
  143. {
  144. name: "无Gemini有Antigravity-返回静态模型信息",
  145. hasGeminiAccount: false,
  146. hasAntigravity: true,
  147. expectedBehavior: "static_model_info",
  148. },
  149. {
  150. name: "无任何账户-返回503",
  151. hasGeminiAccount: false,
  152. hasAntigravity: false,
  153. expectedBehavior: "service_unavailable",
  154. },
  155. }
  156. for _, tt := range tests {
  157. t.Run(tt.name, func(t *testing.T) {
  158. // 模拟 GeminiV1BetaGetModel 的逻辑 (lines 77-87 in gemini_v1beta_handler.go)
  159. var behavior string
  160. if tt.hasGeminiAccount {
  161. behavior = "forward_to_upstream"
  162. } else if tt.hasAntigravity {
  163. behavior = "static_model_info"
  164. } else {
  165. behavior = "service_unavailable"
  166. }
  167. require.Equal(t, tt.expectedBehavior, behavior)
  168. })
  169. }
  170. }
  171. func TestShouldFallbackGeminiModel_KnownFallbackOn404(t *testing.T) {
  172. t.Parallel()
  173. res := &service.UpstreamHTTPResult{StatusCode: http.StatusNotFound}
  174. require.True(t, shouldFallbackGeminiModel("gemini-3.1-pro-preview-customtools", res))
  175. }
  176. func TestShouldFallbackGeminiModel_UnknownModelOn404(t *testing.T) {
  177. t.Parallel()
  178. res := &service.UpstreamHTTPResult{StatusCode: http.StatusNotFound}
  179. require.False(t, shouldFallbackGeminiModel("gemini-future-model", res))
  180. }
  181. func TestShouldFallbackGeminiModel_DelegatesScopeFallback(t *testing.T) {
  182. t.Parallel()
  183. res := &service.UpstreamHTTPResult{
  184. StatusCode: http.StatusForbidden,
  185. Headers: http.Header{"Www-Authenticate": []string{"Bearer error=\"insufficient_scope\""}},
  186. Body: []byte("insufficient authentication scopes"),
  187. }
  188. require.True(t, shouldFallbackGeminiModel("gemini-future-model", res))
  189. }