antigravity_model_mapping_test.go 9.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327
  1. //go:build unit
  2. package service
  3. import (
  4. "testing"
  5. "github.com/stretchr/testify/require"
  6. )
  7. func TestAntigravityGatewayService_GetMappedModel(t *testing.T) {
  8. svc := &AntigravityGatewayService{}
  9. tests := []struct {
  10. name string
  11. requestedModel string
  12. accountMapping map[string]string
  13. expected string
  14. }{
  15. // 1. 账户级映射优先
  16. {
  17. name: "账户映射优先",
  18. requestedModel: "claude-3-5-sonnet-20241022",
  19. accountMapping: map[string]string{"claude-3-5-sonnet-20241022": "custom-model"},
  20. expected: "custom-model",
  21. },
  22. {
  23. name: "账户映射 - 可覆盖默认映射的模型",
  24. requestedModel: "claude-sonnet-4-5",
  25. accountMapping: map[string]string{"claude-sonnet-4-5": "my-custom-sonnet"},
  26. expected: "my-custom-sonnet",
  27. },
  28. {
  29. name: "账户映射 - 可覆盖未知模型",
  30. requestedModel: "claude-opus-4",
  31. accountMapping: map[string]string{"claude-opus-4": "my-opus"},
  32. expected: "my-opus",
  33. },
  34. // 2. 默认映射(DefaultAntigravityModelMapping)
  35. {
  36. name: "默认映射 - claude-opus-4-6 → claude-opus-4-6-thinking",
  37. requestedModel: "claude-opus-4-6",
  38. accountMapping: nil,
  39. expected: "claude-opus-4-6-thinking",
  40. },
  41. {
  42. name: "默认映射 - claude-opus-4-5-20251101 → claude-opus-4-6-thinking",
  43. requestedModel: "claude-opus-4-5-20251101",
  44. accountMapping: nil,
  45. expected: "claude-opus-4-6-thinking",
  46. },
  47. {
  48. name: "默认映射 - claude-opus-4-5-thinking → claude-opus-4-6-thinking",
  49. requestedModel: "claude-opus-4-5-thinking",
  50. accountMapping: nil,
  51. expected: "claude-opus-4-6-thinking",
  52. },
  53. {
  54. name: "默认映射 - claude-haiku-4-5 → claude-sonnet-4-6",
  55. requestedModel: "claude-haiku-4-5",
  56. accountMapping: nil,
  57. expected: "claude-sonnet-4-6",
  58. },
  59. {
  60. name: "默认映射 - claude-haiku-4-5-20251001 → claude-sonnet-4-6",
  61. requestedModel: "claude-haiku-4-5-20251001",
  62. accountMapping: nil,
  63. expected: "claude-sonnet-4-6",
  64. },
  65. {
  66. name: "默认映射 - claude-sonnet-4-5-20250929 → claude-sonnet-4-6",
  67. requestedModel: "claude-sonnet-4-5-20250929",
  68. accountMapping: nil,
  69. expected: "claude-sonnet-4-6",
  70. },
  71. // 3. 默认映射中的透传(映射到自己)
  72. {
  73. name: "默认映射透传 - claude-fable-5-1",
  74. requestedModel: "claude-fable-5-1",
  75. accountMapping: nil,
  76. expected: "claude-fable-5-1",
  77. },
  78. {
  79. name: "默认映射透传 - claude-fable-5",
  80. requestedModel: "claude-fable-5",
  81. accountMapping: nil,
  82. expected: "claude-fable-5",
  83. },
  84. {
  85. name: "默认映射透传 - claude-sonnet-4-6",
  86. requestedModel: "claude-sonnet-4-6",
  87. accountMapping: nil,
  88. expected: "claude-sonnet-4-6",
  89. },
  90. {
  91. name: "显式 canonical 选择 - claude-sonnet-4-5 透传",
  92. requestedModel: "claude-sonnet-4-5",
  93. accountMapping: nil,
  94. expected: "claude-sonnet-4-5",
  95. },
  96. {
  97. name: "默认映射透传 - claude-opus-4-8",
  98. requestedModel: "claude-opus-4-8",
  99. accountMapping: nil,
  100. expected: "claude-opus-4-8",
  101. },
  102. {
  103. name: "默认映射透传 - claude-opus-4-7",
  104. requestedModel: "claude-opus-4-7",
  105. accountMapping: nil,
  106. expected: "claude-opus-4-7",
  107. },
  108. {
  109. name: "默认映射透传 - claude-opus-4-6-thinking",
  110. requestedModel: "claude-opus-4-6-thinking",
  111. accountMapping: nil,
  112. expected: "claude-opus-4-6-thinking",
  113. },
  114. {
  115. name: "默认映射 - claude-sonnet-4-5-thinking → claude-sonnet-4-6",
  116. requestedModel: "claude-sonnet-4-5-thinking",
  117. accountMapping: nil,
  118. expected: "claude-sonnet-4-6",
  119. },
  120. {
  121. name: "账户显式目标只映射一步 - custom-sonnet → claude-sonnet-4-5",
  122. requestedModel: "custom-sonnet",
  123. accountMapping: map[string]string{
  124. "custom-sonnet": "claude-sonnet-4-5",
  125. "claude-sonnet-4-5": "claude-sonnet-4-6",
  126. },
  127. expected: "claude-sonnet-4-5",
  128. },
  129. {
  130. name: "默认映射透传 - gemini-2.5-flash",
  131. requestedModel: "gemini-2.5-flash",
  132. accountMapping: nil,
  133. expected: "gemini-2.5-flash",
  134. },
  135. {
  136. name: "默认映射透传 - gemini-2.5-pro",
  137. requestedModel: "gemini-2.5-pro",
  138. accountMapping: nil,
  139. expected: "gemini-2.5-pro",
  140. },
  141. {
  142. name: "默认映射透传 - gemini-3-flash",
  143. requestedModel: "gemini-3-flash",
  144. accountMapping: nil,
  145. expected: "gemini-3-flash",
  146. },
  147. // 4. 未在默认映射中的模型返回空字符串(不支持)
  148. {
  149. name: "未知模型 - claude-unknown 返回空",
  150. requestedModel: "claude-unknown",
  151. accountMapping: nil,
  152. expected: "",
  153. },
  154. {
  155. name: "未知模型 - claude-3-5-sonnet-20241022 返回空(未在默认映射)",
  156. requestedModel: "claude-3-5-sonnet-20241022",
  157. accountMapping: nil,
  158. expected: "",
  159. },
  160. {
  161. name: "未知模型 - claude-3-opus-20240229 返回空",
  162. requestedModel: "claude-3-opus-20240229",
  163. accountMapping: nil,
  164. expected: "",
  165. },
  166. {
  167. name: "未知模型 - claude-opus-4 返回空",
  168. requestedModel: "claude-opus-4",
  169. accountMapping: nil,
  170. expected: "",
  171. },
  172. {
  173. name: "未知模型 - gemini-future-model 返回空",
  174. requestedModel: "gemini-future-model",
  175. accountMapping: nil,
  176. expected: "",
  177. },
  178. }
  179. for _, tt := range tests {
  180. t.Run(tt.name, func(t *testing.T) {
  181. account := &Account{
  182. Platform: PlatformAntigravity,
  183. }
  184. if tt.accountMapping != nil {
  185. // GetModelMapping 期望 model_mapping 是 map[string]any 格式
  186. mappingAny := make(map[string]any)
  187. for k, v := range tt.accountMapping {
  188. mappingAny[k] = v
  189. }
  190. account.Credentials = map[string]any{
  191. "model_mapping": mappingAny,
  192. }
  193. }
  194. got := svc.getMappedModel(account, tt.requestedModel)
  195. require.Equal(t, tt.expected, got, "model: %s", tt.requestedModel)
  196. })
  197. }
  198. }
  199. func TestAntigravityGatewayService_GetMappedModel_EdgeCases(t *testing.T) {
  200. svc := &AntigravityGatewayService{}
  201. tests := []struct {
  202. name string
  203. requestedModel string
  204. expected string
  205. }{
  206. // 空字符串和非 claude/gemini 前缀返回空字符串
  207. {"空字符串", "", ""},
  208. {"非claude/gemini前缀 - gpt", "gpt-4", ""},
  209. {"非claude/gemini前缀 - llama", "llama-3", ""},
  210. }
  211. for _, tt := range tests {
  212. t.Run(tt.name, func(t *testing.T) {
  213. account := &Account{Platform: PlatformAntigravity}
  214. got := svc.getMappedModel(account, tt.requestedModel)
  215. require.Equal(t, tt.expected, got)
  216. })
  217. }
  218. }
  219. func TestAntigravityGatewayService_IsModelSupported(t *testing.T) {
  220. svc := &AntigravityGatewayService{}
  221. tests := []struct {
  222. name string
  223. model string
  224. expected bool
  225. }{
  226. // 直接支持
  227. {"直接支持 - claude-fable-5-1", "claude-fable-5-1", true},
  228. {"直接支持 - claude-fable-5", "claude-fable-5", true},
  229. {"直接支持 - claude-sonnet-4-5", "claude-sonnet-4-5", true},
  230. {"直接支持 - gemini-3-flash", "gemini-3-flash", true},
  231. // 可映射(有明确前缀映射)
  232. {"可映射 - claude-opus-4-8", "claude-opus-4-8", true},
  233. {"可映射 - claude-opus-4-6", "claude-opus-4-6", true},
  234. // 前缀透传(claude 和 gemini 前缀)
  235. {"Gemini前缀", "gemini-unknown", true},
  236. {"Claude前缀", "claude-unknown", true},
  237. // 不支持
  238. {"不支持 - gpt-4", "gpt-4", false},
  239. {"不支持 - 空字符串", "", false},
  240. }
  241. for _, tt := range tests {
  242. t.Run(tt.name, func(t *testing.T) {
  243. got := svc.IsModelSupported(tt.model)
  244. require.Equal(t, tt.expected, got)
  245. })
  246. }
  247. }
  248. // TestMapAntigravityModel_WildcardTargetEqualsRequest 测试通配符映射目标恰好等于请求模型名的 edge case
  249. // 例如 {"claude-*": "claude-sonnet-4-5"},请求 "claude-sonnet-4-5" 时应该通过
  250. func TestMapAntigravityModel_WildcardTargetEqualsRequest(t *testing.T) {
  251. tests := []struct {
  252. name string
  253. modelMapping map[string]any
  254. requestedModel string
  255. expected string
  256. }{
  257. {
  258. name: "wildcard target equals request model",
  259. modelMapping: map[string]any{"claude-*": "claude-sonnet-4-5"},
  260. requestedModel: "claude-sonnet-4-5",
  261. expected: "claude-sonnet-4-5",
  262. },
  263. {
  264. name: "wildcard target differs from request model",
  265. modelMapping: map[string]any{"claude-*": "claude-sonnet-4-5"},
  266. requestedModel: "claude-opus-4-6",
  267. expected: "claude-sonnet-4-5",
  268. },
  269. {
  270. name: "wildcard no match",
  271. modelMapping: map[string]any{"claude-*": "claude-sonnet-4-5"},
  272. requestedModel: "gpt-4o",
  273. expected: "",
  274. },
  275. {
  276. name: "explicit passthrough same name",
  277. modelMapping: map[string]any{"claude-sonnet-4-5": "claude-sonnet-4-5"},
  278. requestedModel: "claude-sonnet-4-5",
  279. expected: "claude-sonnet-4-5",
  280. },
  281. {
  282. name: "multiple wildcards target equals one request",
  283. modelMapping: map[string]any{"claude-*": "claude-sonnet-4-5", "gemini-*": "gemini-2.5-flash"},
  284. requestedModel: "gemini-2.5-flash",
  285. expected: "gemini-2.5-flash",
  286. },
  287. {
  288. name: "customtools alias falls back to normalized preview mapping",
  289. modelMapping: map[string]any{"gemini-3.1-pro-preview": "gemini-3.1-pro-high"},
  290. requestedModel: "gemini-3.1-pro-preview-customtools",
  291. expected: "gemini-3.1-pro-high",
  292. },
  293. }
  294. for _, tt := range tests {
  295. t.Run(tt.name, func(t *testing.T) {
  296. account := &Account{
  297. Platform: PlatformAntigravity,
  298. Credentials: map[string]any{
  299. "model_mapping": tt.modelMapping,
  300. },
  301. }
  302. got := mapAntigravityModel(account, tt.requestedModel)
  303. require.Equal(t, tt.expected, got, "mapAntigravityModel(%q) = %q, want %q", tt.requestedModel, got, tt.expected)
  304. })
  305. }
  306. }