account_test_service_grok_test.go 25 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644
  1. //go:build unit
  2. package service
  3. import (
  4. "bytes"
  5. "context"
  6. "encoding/base64"
  7. "errors"
  8. "image"
  9. "image/color"
  10. "image/png"
  11. "io"
  12. "net/http"
  13. "net/http/httptest"
  14. "strings"
  15. "testing"
  16. "time"
  17. "github.com/gin-gonic/gin"
  18. "github.com/stretchr/testify/require"
  19. "github.com/tidwall/gjson"
  20. )
  21. type grokAccountTestRateLimitRepo struct {
  22. *mockAccountRepoForGemini
  23. rateLimitedCalls int
  24. resetAt time.Time
  25. }
  26. func TestObserveGrokTestResponseClassifiesBodyOnlyQuotaErrors(t *testing.T) {
  27. account := &Account{ID: 1901, Platform: PlatformGrok, Type: AccountTypeOAuth}
  28. repo := &grokQuotaAccountRepo{mockAccountRepoForPlatform: &mockAccountRepoForPlatform{
  29. accountsByID: map[int64]*Account{account.ID: account},
  30. }}
  31. svc := &AccountTestService{accountRepo: repo}
  32. resp := &http.Response{
  33. StatusCode: http.StatusBadRequest,
  34. Header: make(http.Header),
  35. Body: io.NopCloser(strings.NewReader(`{"error":{"code":"subscription:free-usage-exhausted","message":"included free usage exhausted"}}`)),
  36. }
  37. svc.observeGrokTestResponse(context.Background(), account, resp)
  38. require.Equal(t, 1, repo.tempUnschedCalls)
  39. require.Equal(t, "grok free usage exhausted", repo.lastTempUnschedReason)
  40. body, err := io.ReadAll(resp.Body)
  41. require.NoError(t, err)
  42. require.Contains(t, string(body), "free-usage-exhausted")
  43. }
  44. func TestObserveGrokTestResponseDoesNotQuarantineContentPolicy(t *testing.T) {
  45. account := &Account{ID: 1902, Platform: PlatformGrok, Type: AccountTypeOAuth}
  46. repo := &grokQuotaAccountRepo{mockAccountRepoForPlatform: &mockAccountRepoForPlatform{
  47. accountsByID: map[int64]*Account{account.ID: account},
  48. }}
  49. svc := &AccountTestService{accountRepo: repo}
  50. resp := &http.Response{
  51. StatusCode: http.StatusForbidden,
  52. Header: make(http.Header),
  53. Body: io.NopCloser(strings.NewReader(`{"error":{"code":"new_sensitive","message":"text is sensitive"}}`)),
  54. }
  55. svc.observeGrokTestResponse(context.Background(), account, resp)
  56. require.Zero(t, repo.tempUnschedCalls)
  57. require.Zero(t, repo.rateLimitedCalls)
  58. }
  59. func TestObserveGrokTestResponseKeepsEntitlement403Cooldown(t *testing.T) {
  60. account := &Account{ID: 1903, Platform: PlatformGrok, Type: AccountTypeOAuth}
  61. repo := &grokQuotaAccountRepo{mockAccountRepoForPlatform: &mockAccountRepoForPlatform{
  62. accountsByID: map[int64]*Account{account.ID: account},
  63. }}
  64. svc := &AccountTestService{accountRepo: repo}
  65. resp := &http.Response{
  66. StatusCode: http.StatusForbidden,
  67. Header: make(http.Header),
  68. Body: io.NopCloser(strings.NewReader(`{"error":{"message":"subscription required"}}`)),
  69. }
  70. before := time.Now()
  71. svc.observeGrokTestResponse(context.Background(), account, resp)
  72. require.Equal(t, 1, repo.tempUnschedCalls)
  73. require.Equal(t, "grok entitlement or subscription tier denied", repo.lastTempUnschedReason)
  74. require.Greater(t, repo.lastTempUnschedUntil, before.Add(29*time.Minute))
  75. }
  76. func (r *grokAccountTestRateLimitRepo) SetRateLimited(_ context.Context, _ int64, resetAt time.Time) error {
  77. r.rateLimitedCalls++
  78. r.resetAt = resetAt
  79. return nil
  80. }
  81. func TestAccountTestService_TestAccountConnection_GrokUsesXAIResponses(t *testing.T) {
  82. gin.SetMode(gin.TestMode)
  83. account := &Account{
  84. ID: 13,
  85. Name: "grok-oauth",
  86. Platform: PlatformGrok,
  87. Type: AccountTypeOAuth,
  88. Status: StatusActive,
  89. Schedulable: true,
  90. Concurrency: 1,
  91. Credentials: map[string]any{
  92. "access_token": "grok-access-token",
  93. "refresh_token": "grok-refresh-token",
  94. "expires_at": time.Now().Add(2 * time.Hour).UTC().Format(time.RFC3339),
  95. "model_mapping": map[string]any{
  96. "grok": "grok-4.3",
  97. },
  98. },
  99. }
  100. repo := &mockAccountRepoForGemini{
  101. accountsByID: map[int64]*Account{account.ID: account},
  102. }
  103. upstream := &httpUpstreamRecorder{resp: &http.Response{
  104. StatusCode: http.StatusOK,
  105. Header: http.Header{"Content-Type": []string{"text/event-stream"}},
  106. Body: io.NopCloser(strings.NewReader(
  107. "data: {\"type\":\"response.output_text.delta\",\"delta\":\"ok\"}\n\n" +
  108. "data: {\"type\":\"response.completed\"}\n\n",
  109. )),
  110. }}
  111. svc := &AccountTestService{
  112. accountRepo: repo,
  113. grokTokenProvider: NewGrokTokenProvider(repo, nil),
  114. httpUpstream: upstream,
  115. }
  116. rec := httptest.NewRecorder()
  117. c, _ := gin.CreateTestContext(rec)
  118. c.Request = httptest.NewRequest(http.MethodPost, "/api/v1/admin/accounts/13/test", nil)
  119. err := svc.TestAccountConnection(c, account.ID, "grok", "", AccountTestModeDefault)
  120. require.NoError(t, err)
  121. require.Equal(t, "https://cli-chat-proxy.grok.com/v1/responses", upstream.lastReq.URL.String())
  122. require.Equal(t, "Bearer grok-access-token", upstream.lastReq.Header.Get("Authorization"))
  123. require.Equal(t, grokCLIVersion, upstream.lastReq.Header.Get("X-Grok-Client-Version"))
  124. require.Equal(t, "application/json, text/event-stream", upstream.lastReq.Header.Get("Accept"))
  125. require.Equal(t, "grok-4.3", gjson.GetBytes(upstream.lastBody, "model").String())
  126. require.Equal(t, grokQuotaProbeInput, gjson.GetBytes(upstream.lastBody, "input").String())
  127. require.True(t, gjson.GetBytes(upstream.lastBody, "stream").Bool())
  128. require.False(t, gjson.GetBytes(upstream.lastBody, "max_output_tokens").Exists())
  129. require.False(t, gjson.GetBytes(upstream.lastBody, "store").Exists())
  130. require.NotContains(t, rec.Body.String(), "claude")
  131. require.Contains(t, rec.Body.String(), `"model":"grok-4.3"`)
  132. require.Contains(t, rec.Body.String(), `"type":"test_complete"`)
  133. }
  134. func TestAccountTestService_TestAccountConnection_GrokDefaultsEmptyModelTo45(t *testing.T) {
  135. gin.SetMode(gin.TestMode)
  136. account := &Account{
  137. ID: 16,
  138. Name: "grok-oauth-default-model",
  139. Platform: PlatformGrok,
  140. Type: AccountTypeOAuth,
  141. Status: StatusActive,
  142. Schedulable: true,
  143. Concurrency: 1,
  144. Credentials: map[string]any{
  145. "access_token": "grok-access-token",
  146. "refresh_token": "grok-refresh-token",
  147. "expires_at": time.Now().Add(2 * time.Hour).UTC().Format(time.RFC3339),
  148. },
  149. }
  150. repo := &mockAccountRepoForGemini{accountsByID: map[int64]*Account{account.ID: account}}
  151. upstream := &httpUpstreamRecorder{resp: &http.Response{
  152. StatusCode: http.StatusOK,
  153. Header: http.Header{"Content-Type": []string{"text/event-stream"}},
  154. Body: io.NopCloser(strings.NewReader(
  155. "data: {\"type\":\"response.output_text.delta\",\"delta\":\"ok\"}\n\n" +
  156. "data: {\"type\":\"response.completed\"}\n\n",
  157. )),
  158. }}
  159. svc := &AccountTestService{
  160. accountRepo: repo,
  161. grokTokenProvider: NewGrokTokenProvider(repo, nil),
  162. httpUpstream: upstream,
  163. }
  164. recorder := httptest.NewRecorder()
  165. c, _ := gin.CreateTestContext(recorder)
  166. c.Request = httptest.NewRequest(http.MethodPost, "/api/v1/admin/accounts/16/test", nil)
  167. err := svc.TestAccountConnection(c, account.ID, "", "", AccountTestModeDefault)
  168. require.NoError(t, err)
  169. require.Equal(t, grokDefaultResponsesModel, gjson.GetBytes(upstream.lastBody, "model").String())
  170. require.Contains(t, recorder.Body.String(), `"model":"grok-4.5"`)
  171. }
  172. func TestAccountTestService_Grok429PersistsRateLimitReset(t *testing.T) {
  173. gin.SetMode(gin.TestMode)
  174. account := &Account{
  175. ID: 14,
  176. Name: "grok-oauth-limited",
  177. Platform: PlatformGrok,
  178. Type: AccountTypeOAuth,
  179. Status: StatusActive,
  180. Schedulable: true,
  181. Concurrency: 1,
  182. Credentials: map[string]any{
  183. "access_token": "grok-access-token",
  184. "refresh_token": "grok-refresh-token",
  185. "expires_at": time.Now().Add(2 * time.Hour).UTC().Format(time.RFC3339),
  186. },
  187. }
  188. baseRepo := &mockAccountRepoForGemini{accountsByID: map[int64]*Account{account.ID: account}}
  189. repo := &grokAccountTestRateLimitRepo{mockAccountRepoForGemini: baseRepo}
  190. upstream := &httpUpstreamRecorder{resp: &http.Response{
  191. StatusCode: http.StatusTooManyRequests,
  192. Header: http.Header{"Retry-After": []string{"45"}},
  193. Body: io.NopCloser(strings.NewReader(`{"error":{"message":"rate limited"}}`)),
  194. }}
  195. svc := &AccountTestService{
  196. accountRepo: repo,
  197. grokTokenProvider: NewGrokTokenProvider(repo, nil),
  198. httpUpstream: upstream,
  199. }
  200. recorder := httptest.NewRecorder()
  201. c, _ := gin.CreateTestContext(recorder)
  202. c.Request = httptest.NewRequest(http.MethodPost, "/api/v1/admin/accounts/14/test", nil)
  203. err := svc.TestAccountConnection(c, account.ID, "grok", "", AccountTestModeDefault)
  204. require.Error(t, err)
  205. require.Equal(t, 1, repo.rateLimitedCalls)
  206. require.WithinDuration(t, time.Now().Add(45*time.Second), repo.resetAt, time.Second)
  207. }
  208. func TestAccountTestService_Grok429WithoutQuotaHeadersUsesFallback(t *testing.T) {
  209. gin.SetMode(gin.TestMode)
  210. account := &Account{
  211. ID: 15, Name: "grok-oauth-limited-no-headers", Platform: PlatformGrok,
  212. Type: AccountTypeOAuth, Status: StatusActive, Schedulable: true, Concurrency: 1,
  213. Credentials: map[string]any{
  214. "access_token": "grok-access-token",
  215. "refresh_token": "grok-refresh-token",
  216. "expires_at": time.Now().Add(2 * time.Hour).UTC().Format(time.RFC3339),
  217. },
  218. }
  219. baseRepo := &mockAccountRepoForGemini{accountsByID: map[int64]*Account{account.ID: account}}
  220. repo := &grokAccountTestRateLimitRepo{mockAccountRepoForGemini: baseRepo}
  221. upstream := &httpUpstreamRecorder{resp: &http.Response{
  222. StatusCode: http.StatusTooManyRequests,
  223. Body: io.NopCloser(strings.NewReader(`{"error":{"message":"quota exhausted"}}`)),
  224. }}
  225. svc := &AccountTestService{
  226. accountRepo: repo, grokTokenProvider: NewGrokTokenProvider(repo, nil), httpUpstream: upstream,
  227. }
  228. recorder := httptest.NewRecorder()
  229. c, _ := gin.CreateTestContext(recorder)
  230. c.Request = httptest.NewRequest(http.MethodPost, "/api/v1/admin/accounts/15/test", nil)
  231. before := time.Now()
  232. err := svc.TestAccountConnection(c, account.ID, "grok", "", AccountTestModeDefault)
  233. require.Error(t, err)
  234. require.Equal(t, 1, repo.rateLimitedCalls)
  235. require.WithinDuration(t, before.Add(grokRateLimitFallbackCooldown), repo.resetAt, time.Second)
  236. }
  237. func TestAccountTestService_GrokImageModelUsesImagesGenerations(t *testing.T) {
  238. gin.SetMode(gin.TestMode)
  239. account := &Account{
  240. ID: 17, Name: "grok-oauth-image", Platform: PlatformGrok,
  241. Type: AccountTypeOAuth, Status: StatusActive, Schedulable: true, Concurrency: 1,
  242. Credentials: map[string]any{
  243. "access_token": "grok-access-token",
  244. "refresh_token": "grok-refresh-token",
  245. "expires_at": time.Now().Add(2 * time.Hour).UTC().Format(time.RFC3339),
  246. "base_url": "https://cli-chat-proxy.grok.com/v1",
  247. },
  248. }
  249. repo := &mockAccountRepoForGemini{accountsByID: map[int64]*Account{account.ID: account}}
  250. upstream := &httpUpstreamRecorder{resp: &http.Response{
  251. StatusCode: http.StatusOK,
  252. Header: http.Header{"Content-Type": []string{"application/json"}},
  253. Body: io.NopCloser(strings.NewReader(
  254. `{"data":[{"b64_json":"QUJD","mime_type":"image/jpeg"}]}`,
  255. )),
  256. }}
  257. svc := &AccountTestService{
  258. accountRepo: repo,
  259. grokTokenProvider: NewGrokTokenProvider(repo, nil),
  260. httpUpstream: upstream,
  261. }
  262. rec := httptest.NewRecorder()
  263. c, _ := gin.CreateTestContext(rec)
  264. c.Request = httptest.NewRequest(http.MethodPost, "/api/v1/admin/accounts/17/test", nil)
  265. err := svc.TestAccountConnection(c, account.ID, "grok-imagine-image", "a red apple", AccountTestModeDefault)
  266. require.NoError(t, err)
  267. require.Equal(t, "https://api.x.ai/v1/images/generations", upstream.lastReq.URL.String())
  268. require.Equal(t, "grok-imagine-image", gjson.GetBytes(upstream.lastBody, "model").String())
  269. require.Equal(t, "a red apple", gjson.GetBytes(upstream.lastBody, "prompt").String())
  270. require.Equal(t, "b64_json", gjson.GetBytes(upstream.lastBody, "response_format").String())
  271. require.Contains(t, rec.Body.String(), `"type":"image"`)
  272. require.Contains(t, rec.Body.String(), "data:image/jpeg;base64,QUJD")
  273. require.Contains(t, rec.Body.String(), `"type":"test_complete"`)
  274. }
  275. func TestAccountTestService_GrokWebSearchModeUsesResponsesWebSearchTool(t *testing.T) {
  276. gin.SetMode(gin.TestMode)
  277. account := &Account{
  278. ID: 18, Name: "grok-oauth-search", Platform: PlatformGrok,
  279. Type: AccountTypeOAuth, Status: StatusActive, Schedulable: true, Concurrency: 1,
  280. Credentials: map[string]any{
  281. "access_token": "grok-access-token",
  282. "refresh_token": "grok-refresh-token",
  283. "expires_at": time.Now().Add(2 * time.Hour).UTC().Format(time.RFC3339),
  284. },
  285. }
  286. repo := &mockAccountRepoForGemini{accountsByID: map[int64]*Account{account.ID: account}}
  287. upstream := &httpUpstreamRecorder{resp: &http.Response{
  288. StatusCode: http.StatusOK,
  289. Body: io.NopCloser(strings.NewReader(
  290. `{"id":"r1","output":[{"type":"web_search_call","id":"ws1","status":"completed"},{"type":"message","role":"assistant","content":[{"type":"output_text","text":"Grok is built by xAI."}]}]}`,
  291. )),
  292. }}
  293. svc := &AccountTestService{
  294. accountRepo: repo,
  295. grokTokenProvider: NewGrokTokenProvider(repo, nil),
  296. httpUpstream: upstream,
  297. }
  298. rec := httptest.NewRecorder()
  299. c, _ := gin.CreateTestContext(rec)
  300. c.Request = httptest.NewRequest(http.MethodPost, "/api/v1/admin/accounts/18/test", nil)
  301. err := svc.TestAccountConnection(c, account.ID, "grok-4.5", "xAI Grok", AccountTestModeGrokSearch)
  302. require.NoError(t, err)
  303. require.Equal(t, "https://cli-chat-proxy.grok.com/v1/responses", upstream.lastReq.URL.String())
  304. // Standalone web_search wraps the query in the gateway-style prompt.
  305. require.Contains(t, gjson.GetBytes(upstream.lastBody, "input").String(), "xAI Grok")
  306. require.Equal(t, "web_search", gjson.GetBytes(upstream.lastBody, "tools.0.type").String())
  307. require.Equal(t, "web_search_call.action.sources", gjson.GetBytes(upstream.lastBody, "include.0").String())
  308. require.Contains(t, rec.Body.String(), "web_search ok")
  309. require.Contains(t, rec.Body.String(), `"type":"test_complete"`)
  310. }
  311. func TestAccountTestService_GrokTTSIncludesLanguage(t *testing.T) {
  312. gin.SetMode(gin.TestMode)
  313. account := &Account{
  314. ID: 19, Name: "grok-oauth-tts", Platform: PlatformGrok,
  315. Type: AccountTypeOAuth, Status: StatusActive, Schedulable: true, Concurrency: 1,
  316. Credentials: map[string]any{
  317. "access_token": "grok-access-token",
  318. "refresh_token": "grok-refresh-token",
  319. "expires_at": time.Now().Add(2 * time.Hour).UTC().Format(time.RFC3339),
  320. },
  321. }
  322. repo := &mockAccountRepoForGemini{accountsByID: map[int64]*Account{account.ID: account}}
  323. upstream := &httpUpstreamRecorder{resp: &http.Response{
  324. StatusCode: http.StatusOK,
  325. Header: http.Header{"Content-Type": []string{"audio/mpeg"}},
  326. Body: io.NopCloser(strings.NewReader("ID3fakeaudio")),
  327. }}
  328. svc := &AccountTestService{
  329. accountRepo: repo,
  330. grokTokenProvider: NewGrokTokenProvider(repo, nil),
  331. httpUpstream: upstream,
  332. }
  333. rec := httptest.NewRecorder()
  334. c, _ := gin.CreateTestContext(rec)
  335. c.Request = httptest.NewRequest(http.MethodPost, "/api/v1/admin/accounts/19/test", nil)
  336. err := svc.TestAccountConnection(c, account.ID, "", "hello voice", AccountTestModeGrokTTS)
  337. require.NoError(t, err)
  338. require.Equal(t, "https://api.x.ai/v1/tts", upstream.lastReq.URL.String())
  339. require.Equal(t, "hello voice", gjson.GetBytes(upstream.lastBody, "text").String())
  340. require.Equal(t, "en", gjson.GetBytes(upstream.lastBody, "language").String())
  341. require.Contains(t, rec.Body.String(), "tts ok")
  342. require.Contains(t, rec.Body.String(), `"type":"audio"`)
  343. require.Contains(t, rec.Body.String(), "data:audio/mpeg;base64,")
  344. require.Contains(t, rec.Body.String(), `"type":"test_complete"`)
  345. }
  346. func TestAccountTestService_GrokImageEditUsesUploadedImage(t *testing.T) {
  347. gin.SetMode(gin.TestMode)
  348. account := &Account{
  349. ID: 24, Name: "grok-oauth-image-edit", Platform: PlatformGrok,
  350. Type: AccountTypeOAuth, Status: StatusActive, Schedulable: true, Concurrency: 1,
  351. Credentials: map[string]any{
  352. "access_token": "grok-access-token",
  353. "refresh_token": "grok-refresh-token",
  354. "expires_at": time.Now().Add(2 * time.Hour).UTC().Format(time.RFC3339),
  355. },
  356. }
  357. repo := &mockAccountRepoForGemini{accountsByID: map[int64]*Account{account.ID: account}}
  358. upstream := &httpUpstreamRecorder{resp: &http.Response{
  359. StatusCode: http.StatusOK,
  360. Header: http.Header{"Content-Type": []string{"application/json"}},
  361. Body: io.NopCloser(strings.NewReader(
  362. `{"data":[{"b64_json":"QUJD","mime_type":"image/png"}]}`,
  363. )),
  364. }}
  365. svc := &AccountTestService{
  366. accountRepo: repo,
  367. grokTokenProvider: NewGrokTokenProvider(repo, nil),
  368. httpUpstream: upstream,
  369. }
  370. rec := httptest.NewRecorder()
  371. c, _ := gin.CreateTestContext(rec)
  372. c.Request = httptest.NewRequest(http.MethodPost, "/api/v1/admin/accounts/24/test", nil)
  373. // 8x8 solid PNG (xAI min dimension is 8px).
  374. src := minimalAccountTestPNGDataURL(8, 8)
  375. err := svc.TestAccountConnection(c, account.ID, "grok-imagine-image", "edit me", AccountTestModeGrokImage, AccountTestOptions{
  376. ImageDataURL: src,
  377. })
  378. require.NoError(t, err)
  379. require.Equal(t, "https://api.x.ai/v1/images/edits", upstream.lastReq.URL.String())
  380. require.True(t, strings.HasPrefix(gjson.GetBytes(upstream.lastBody, "image.url").String(), "data:image/png;base64,"))
  381. require.Equal(t, "image_url", gjson.GetBytes(upstream.lastBody, "image.type").String())
  382. require.Equal(t, "b64_json", gjson.GetBytes(upstream.lastBody, "response_format").String())
  383. // concrete image model ids pass through; only bare "grok-imagine" is aliased.
  384. require.Equal(t, "grok-imagine-image", gjson.GetBytes(upstream.lastBody, "model").String())
  385. require.Contains(t, rec.Body.String(), `"type":"image"`)
  386. }
  387. func TestAccountTestService_GrokImageEditRejectsTinySource(t *testing.T) {
  388. gin.SetMode(gin.TestMode)
  389. account := &Account{
  390. ID: 25, Name: "grok-oauth-image-tiny", Platform: PlatformGrok,
  391. Type: AccountTypeOAuth, Status: StatusActive, Schedulable: true, Concurrency: 1,
  392. Credentials: map[string]any{
  393. "access_token": "grok-access-token",
  394. "refresh_token": "grok-refresh-token",
  395. "expires_at": time.Now().Add(2 * time.Hour).UTC().Format(time.RFC3339),
  396. },
  397. }
  398. repo := &mockAccountRepoForGemini{accountsByID: map[int64]*Account{account.ID: account}}
  399. svc := &AccountTestService{
  400. accountRepo: repo,
  401. grokTokenProvider: NewGrokTokenProvider(repo, nil),
  402. httpUpstream: &httpUpstreamRecorder{},
  403. }
  404. rec := httptest.NewRecorder()
  405. c, _ := gin.CreateTestContext(rec)
  406. c.Request = httptest.NewRequest(http.MethodPost, "/api/v1/admin/accounts/25/test", nil)
  407. // 1x1 PNG data URL
  408. tiny := "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8z8BQDwAEhQGAhKmMIQAAAABJRU5ErkJggg=="
  409. err := svc.TestAccountConnection(c, account.ID, "grok-imagine-image-quality", "edit", AccountTestModeGrokImage, AccountTestOptions{
  410. ImageDataURL: tiny,
  411. })
  412. require.Error(t, err)
  413. require.Contains(t, rec.Body.String(), "too small")
  414. }
  415. // minimalAccountTestPNGDataURL builds a solid RGBA PNG as a data URL for tests.
  416. func minimalAccountTestPNGDataURL(w, h int) string {
  417. img := image.NewRGBA(image.Rect(0, 0, w, h))
  418. for y := 0; y < h; y++ {
  419. for x := 0; x < w; x++ {
  420. img.Set(x, y, color.RGBA{R: 200, G: 40, B: 40, A: 255})
  421. }
  422. }
  423. var buf bytes.Buffer
  424. _ = png.Encode(&buf, img)
  425. return "data:image/png;base64," + base64.StdEncoding.EncodeToString(buf.Bytes())
  426. }
  427. func TestAccountTestService_GrokExplicitImageModeDefaultsModel(t *testing.T) {
  428. gin.SetMode(gin.TestMode)
  429. account := &Account{
  430. ID: 20, Name: "grok-oauth-image-mode", Platform: PlatformGrok,
  431. Type: AccountTypeOAuth, Status: StatusActive, Schedulable: true, Concurrency: 1,
  432. Credentials: map[string]any{
  433. "access_token": "grok-access-token",
  434. "refresh_token": "grok-refresh-token",
  435. "expires_at": time.Now().Add(2 * time.Hour).UTC().Format(time.RFC3339),
  436. },
  437. }
  438. repo := &mockAccountRepoForGemini{accountsByID: map[int64]*Account{account.ID: account}}
  439. upstream := &httpUpstreamRecorder{resp: &http.Response{
  440. StatusCode: http.StatusOK,
  441. Header: http.Header{"Content-Type": []string{"application/json"}},
  442. Body: io.NopCloser(strings.NewReader(
  443. `{"data":[{"b64_json":"QUJD","mime_type":"image/jpeg"}]}`,
  444. )),
  445. }}
  446. svc := &AccountTestService{
  447. accountRepo: repo,
  448. grokTokenProvider: NewGrokTokenProvider(repo, nil),
  449. httpUpstream: upstream,
  450. }
  451. rec := httptest.NewRecorder()
  452. c, _ := gin.CreateTestContext(rec)
  453. c.Request = httptest.NewRequest(http.MethodPost, "/api/v1/admin/accounts/20/test", nil)
  454. err := svc.TestAccountConnection(c, account.ID, "", "", AccountTestModeGrokImage)
  455. require.NoError(t, err)
  456. require.Equal(t, "https://api.x.ai/v1/images/generations", upstream.lastReq.URL.String())
  457. require.Equal(t, "grok-imagine-image", gjson.GetBytes(upstream.lastBody, "model").String())
  458. require.Equal(t, "b64_json", gjson.GetBytes(upstream.lastBody, "response_format").String())
  459. require.Contains(t, rec.Body.String(), `"type":"test_complete"`)
  460. }
  461. func TestAccountTestService_GrokVideoUpstreamErrorIsNotMaskedAsSuccess(t *testing.T) {
  462. gin.SetMode(gin.TestMode)
  463. account := &Account{
  464. ID: 21, Name: "grok-oauth-video-err", Platform: PlatformGrok,
  465. Type: AccountTypeOAuth, Status: StatusActive, Schedulable: true, Concurrency: 1,
  466. Credentials: map[string]any{
  467. "access_token": "grok-access-token",
  468. "refresh_token": "grok-refresh-token",
  469. "expires_at": time.Now().Add(2 * time.Hour).UTC().Format(time.RFC3339),
  470. },
  471. }
  472. repo := &mockAccountRepoForGemini{accountsByID: map[int64]*Account{account.ID: account}}
  473. upstream := &httpUpstreamRecorder{resp: &http.Response{
  474. StatusCode: http.StatusBadRequest,
  475. Header: http.Header{"Content-Type": []string{"application/json"}},
  476. Body: io.NopCloser(strings.NewReader(
  477. `{"code":"invalid-argument","error":"bad video request"}`,
  478. )),
  479. }}
  480. svc := &AccountTestService{
  481. accountRepo: repo,
  482. grokTokenProvider: NewGrokTokenProvider(repo, nil),
  483. httpUpstream: upstream,
  484. }
  485. rec := httptest.NewRecorder()
  486. c, _ := gin.CreateTestContext(rec)
  487. c.Request = httptest.NewRequest(http.MethodPost, "/api/v1/admin/accounts/21/test", nil)
  488. err := svc.TestAccountConnection(c, account.ID, "grok-imagine-video", "bounce ball", AccountTestModeGrokVideo)
  489. require.Error(t, err)
  490. require.Equal(t, "https://api.x.ai/v1/videos/generations", upstream.lastReq.URL.String())
  491. require.Contains(t, rec.Body.String(), `"type":"error"`)
  492. require.Contains(t, rec.Body.String(), "Grok videos API returned 400")
  493. require.NotContains(t, rec.Body.String(), `"success":true`)
  494. }
  495. type grokRealtimeTestConn struct {
  496. msg []byte
  497. }
  498. func (c *grokRealtimeTestConn) WriteJSON(context.Context, any) error { return nil }
  499. func (c *grokRealtimeTestConn) ReadMessage(context.Context) ([]byte, error) {
  500. if c == nil || len(c.msg) == 0 {
  501. return nil, context.DeadlineExceeded
  502. }
  503. return c.msg, nil
  504. }
  505. func (c *grokRealtimeTestConn) Ping(context.Context) error { return nil }
  506. func (c *grokRealtimeTestConn) Close() error { return nil }
  507. type grokRealtimeTestDialer struct {
  508. lastURL string
  509. lastAuth string
  510. lastProxy string
  511. conn openAIWSClientConn
  512. err error
  513. status int
  514. }
  515. func (d *grokRealtimeTestDialer) Dial(_ context.Context, wsURL string, headers http.Header, proxyURL string) (openAIWSClientConn, int, http.Header, error) {
  516. d.lastURL = wsURL
  517. d.lastAuth = headers.Get("Authorization")
  518. d.lastProxy = proxyURL
  519. if d.err != nil {
  520. return nil, d.status, nil, d.err
  521. }
  522. if d.conn == nil {
  523. d.conn = &grokRealtimeTestConn{}
  524. }
  525. return d.conn, 0, nil, nil
  526. }
  527. func TestAccountTestService_GrokRealtimeModeDialsWS(t *testing.T) {
  528. gin.SetMode(gin.TestMode)
  529. account := &Account{
  530. ID: 22, Name: "grok-oauth-realtime", Platform: PlatformGrok,
  531. Type: AccountTypeOAuth, Status: StatusActive, Schedulable: true, Concurrency: 1,
  532. Credentials: map[string]any{
  533. "access_token": "grok-access-token",
  534. "refresh_token": "grok-refresh-token",
  535. "expires_at": time.Now().Add(2 * time.Hour).UTC().Format(time.RFC3339),
  536. },
  537. }
  538. repo := &mockAccountRepoForGemini{accountsByID: map[int64]*Account{account.ID: account}}
  539. dialer := &grokRealtimeTestDialer{
  540. conn: &grokRealtimeTestConn{msg: []byte(`{"type":"session.created","session":{"id":"sess_1"}}`)},
  541. }
  542. svc := &AccountTestService{
  543. accountRepo: repo,
  544. grokTokenProvider: NewGrokTokenProvider(repo, nil),
  545. grokWSDialer: dialer,
  546. }
  547. rec := httptest.NewRecorder()
  548. c, _ := gin.CreateTestContext(rec)
  549. c.Request = httptest.NewRequest(http.MethodPost, "/api/v1/admin/accounts/22/test", nil)
  550. err := svc.TestAccountConnection(c, account.ID, "", "", AccountTestModeGrokRealtime)
  551. require.NoError(t, err)
  552. require.Contains(t, dialer.lastURL, "wss://api.x.ai/v1/realtime")
  553. require.Contains(t, dialer.lastURL, "model=grok-voice-latest")
  554. require.Equal(t, "Bearer grok-access-token", dialer.lastAuth)
  555. require.Contains(t, rec.Body.String(), "realtime ws handshake ok")
  556. require.Contains(t, rec.Body.String(), "session.created")
  557. require.Contains(t, rec.Body.String(), `"type":"test_complete"`)
  558. require.Contains(t, rec.Body.String(), `"success":true`)
  559. }
  560. func TestAccountTestService_GrokRealtimeModeDialFailure(t *testing.T) {
  561. gin.SetMode(gin.TestMode)
  562. account := &Account{
  563. ID: 23, Name: "grok-oauth-realtime-fail", Platform: PlatformGrok,
  564. Type: AccountTypeOAuth, Status: StatusActive, Schedulable: true, Concurrency: 1,
  565. Credentials: map[string]any{
  566. "access_token": "grok-access-token",
  567. "refresh_token": "grok-refresh-token",
  568. "expires_at": time.Now().Add(2 * time.Hour).UTC().Format(time.RFC3339),
  569. },
  570. }
  571. repo := &mockAccountRepoForGemini{accountsByID: map[int64]*Account{account.ID: account}}
  572. dialer := &grokRealtimeTestDialer{
  573. status: 401,
  574. err: &openAIWSHandshakeError{Body: []byte(`{"error":"unauthorized"}`), Err: errors.New("websocket handshake failed")},
  575. }
  576. svc := &AccountTestService{
  577. accountRepo: repo,
  578. grokTokenProvider: NewGrokTokenProvider(repo, nil),
  579. grokWSDialer: dialer,
  580. }
  581. rec := httptest.NewRecorder()
  582. c, _ := gin.CreateTestContext(rec)
  583. c.Request = httptest.NewRequest(http.MethodPost, "/api/v1/admin/accounts/23/test", nil)
  584. err := svc.TestAccountConnection(c, account.ID, "", "", AccountTestModeGrokRealtime)
  585. require.Error(t, err)
  586. require.Contains(t, rec.Body.String(), `"type":"error"`)
  587. require.Contains(t, rec.Body.String(), "Realtime")
  588. }