upstream_billing_probe_persistence_integration_test.go 20 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544
  1. //go:build integration
  2. package repository
  3. import (
  4. "context"
  5. "encoding/json"
  6. "testing"
  7. "time"
  8. "github.com/Wei-Shaw/sub2api/internal/service"
  9. "github.com/stretchr/testify/require"
  10. )
  11. func TestAccountUpdatePreservesConcurrentProbeSnapshot(t *testing.T) {
  12. ctx := context.Background()
  13. tx := testEntTx(t)
  14. repo := newAccountRepositoryWithSQL(tx.Client(), tx, nil)
  15. account := mustCreateAccount(t, tx.Client(), &service.Account{
  16. Name: "probe-update-preserve",
  17. Platform: service.PlatformOpenAI,
  18. Type: service.AccountTypeAPIKey,
  19. Credentials: map[string]any{"api_key": "sk-old"},
  20. Extra: map[string]any{service.UpstreamBillingProbeEnabledExtraKey: true},
  21. })
  22. stale, err := repo.GetByID(ctx, account.ID)
  23. require.NoError(t, err)
  24. require.NotContains(t, stale.Extra, service.UpstreamBillingProbeExtraKey)
  25. require.NoError(t, repo.UpdateUpstreamBillingProbeSnapshot(ctx, stale, &service.UpstreamBillingProbeSnapshot{
  26. Status: service.UpstreamBillingProbeStatusOK,
  27. LastAttemptAt: time.Now().UTC(),
  28. }, nil))
  29. stale.Name = "ordinary-edit"
  30. require.NoError(t, repo.Update(ctx, stale))
  31. got, err := repo.GetByID(ctx, account.ID)
  32. require.NoError(t, err)
  33. snapshot, ok := got.Extra[service.UpstreamBillingProbeExtraKey].(map[string]any)
  34. require.True(t, ok)
  35. require.Equal(t, service.UpstreamBillingProbeStatusOK, snapshot["status"])
  36. require.NoError(t, repo.UpdateExtra(ctx, got.ID, map[string]any{service.UpstreamBillingProbeEnabledExtraKey: false}))
  37. disabled, err := repo.GetByID(ctx, account.ID)
  38. require.NoError(t, err)
  39. require.NotContains(t, disabled.Extra, service.UpstreamBillingProbeExtraKey)
  40. }
  41. func TestAdminAccountEditPreservesRateSynchronizedAfterLoad(t *testing.T) {
  42. ctx := context.Background()
  43. tx := testEntTx(t)
  44. repo := newAccountRepositoryWithSQL(tx.Client(), tx, nil)
  45. initialRate := 0.1
  46. account := mustCreateAccount(t, tx.Client(), &service.Account{
  47. Name: "probe-rate-concurrent-edit",
  48. Platform: service.PlatformOpenAI,
  49. Type: service.AccountTypeAPIKey,
  50. RateMultiplier: &initialRate,
  51. Credentials: map[string]any{"api_key": "sk-test"},
  52. Extra: map[string]any{
  53. service.UpstreamBillingProbeEnabledExtraKey: true,
  54. service.UpstreamBillingRateSyncEnabledExtraKey: true,
  55. },
  56. })
  57. staleAdminEdit, err := repo.GetByID(ctx, account.ID)
  58. require.NoError(t, err)
  59. probeAccount, err := repo.GetByID(ctx, account.ID)
  60. require.NoError(t, err)
  61. synchronizedRate := 0.2
  62. require.NoError(t, repo.UpdateUpstreamBillingProbeSnapshot(ctx, probeAccount, &service.UpstreamBillingProbeSnapshot{
  63. Status: service.UpstreamBillingProbeStatusOK,
  64. LastAttemptAt: time.Now().UTC(),
  65. }, &synchronizedRate))
  66. staleAdminEdit.Name = "name-only-edit"
  67. require.NoError(t, repo.UpdateWithAccountBillingSettings(ctx, staleAdminEdit, nil, nil, nil))
  68. got, err := repo.GetByID(ctx, account.ID)
  69. require.NoError(t, err)
  70. require.Equal(t, "name-only-edit", got.Name)
  71. require.NotNil(t, got.RateMultiplier)
  72. require.Equal(t, synchronizedRate, *got.RateMultiplier)
  73. }
  74. func TestProbeSnapshotSyncsRateOnlyForSuccessfulEnabledAccount(t *testing.T) {
  75. ctx := context.Background()
  76. tx := testEntTx(t)
  77. repo := newAccountRepositoryWithSQL(tx.Client(), tx, nil)
  78. initialRate := 0.25
  79. account := mustCreateAccount(t, tx.Client(), &service.Account{
  80. Name: "probe-rate-sync",
  81. Platform: service.PlatformGemini,
  82. Type: service.AccountTypeAPIKey,
  83. RateMultiplier: &initialRate,
  84. Credentials: map[string]any{"api_key": "sk-test"},
  85. Extra: map[string]any{
  86. service.UpstreamBillingProbeEnabledExtraKey: true,
  87. service.UpstreamBillingRateSyncEnabledExtraKey: true,
  88. },
  89. })
  90. loaded, err := repo.GetByID(ctx, account.ID)
  91. require.NoError(t, err)
  92. syncedRate := 0.065
  93. require.NoError(t, repo.UpdateUpstreamBillingProbeSnapshot(ctx, loaded, &service.UpstreamBillingProbeSnapshot{
  94. Status: service.UpstreamBillingProbeStatusOK,
  95. LastAttemptAt: time.Now().UTC(),
  96. }, &syncedRate))
  97. got, err := repo.GetByID(ctx, account.ID)
  98. require.NoError(t, err)
  99. require.NotNil(t, got.RateMultiplier)
  100. require.Equal(t, syncedRate, *got.RateMultiplier)
  101. failedRate := 0.9
  102. require.NoError(t, repo.UpdateUpstreamBillingProbeSnapshot(ctx, got, &service.UpstreamBillingProbeSnapshot{
  103. Status: service.UpstreamBillingProbeStatusFailed,
  104. LastAttemptAt: time.Now().UTC(),
  105. }, &failedRate))
  106. got, err = repo.GetByID(ctx, account.ID)
  107. require.NoError(t, err)
  108. require.NotNil(t, got.RateMultiplier)
  109. require.Equal(t, syncedRate, *got.RateMultiplier)
  110. require.NoError(t, repo.UpdateExtra(ctx, account.ID, map[string]any{
  111. service.UpstreamBillingRateSyncEnabledExtraKey: false,
  112. }))
  113. manual, err := repo.GetByID(ctx, account.ID)
  114. require.NoError(t, err)
  115. manualProbeRate := 0.4
  116. require.NoError(t, repo.UpdateUpstreamBillingProbeSnapshot(ctx, manual, &service.UpstreamBillingProbeSnapshot{
  117. Status: service.UpstreamBillingProbeStatusOK,
  118. LastAttemptAt: time.Now().UTC(),
  119. }, &manualProbeRate))
  120. got, err = repo.GetByID(ctx, account.ID)
  121. require.NoError(t, err)
  122. require.NotNil(t, got.RateMultiplier)
  123. require.Equal(t, syncedRate, *got.RateMultiplier)
  124. require.NoError(t, repo.UpdateExtra(ctx, account.ID, map[string]any{
  125. service.UpstreamBillingProbeEnabledExtraKey: false,
  126. service.UpstreamBillingRateSyncEnabledExtraKey: false,
  127. }))
  128. disabled, err := repo.GetByID(ctx, account.ID)
  129. require.NoError(t, err)
  130. require.NoError(t, repo.UpdateUpstreamBillingProbeSnapshot(ctx, disabled, &service.UpstreamBillingProbeSnapshot{
  131. Status: service.UpstreamBillingProbeStatusOK,
  132. LastAttemptAt: time.Now().UTC(),
  133. }, &manualProbeRate))
  134. got, err = repo.GetByID(ctx, account.ID)
  135. require.NoError(t, err)
  136. require.NotNil(t, got.RateMultiplier)
  137. require.Equal(t, syncedRate, *got.RateMultiplier)
  138. }
  139. func TestAccountUpdatePreservesConcurrentProbeEnableFlag(t *testing.T) {
  140. ctx := context.Background()
  141. tx := testEntTx(t)
  142. repo := newAccountRepositoryWithSQL(tx.Client(), tx, nil)
  143. account := mustCreateAccount(t, tx.Client(), &service.Account{
  144. Name: "probe-update-enable",
  145. Platform: service.PlatformOpenAI,
  146. Type: service.AccountTypeAPIKey,
  147. Credentials: map[string]any{"api_key": "sk-test"},
  148. Extra: map[string]any{
  149. service.UpstreamBillingProbeEnabledExtraKey: true,
  150. service.UpstreamBillingProbeExtraKey: map[string]any{"status": service.UpstreamBillingProbeStatusOK},
  151. },
  152. })
  153. stale, err := repo.GetByID(ctx, account.ID)
  154. require.NoError(t, err)
  155. require.NoError(t, repo.UpdateExtra(ctx, account.ID, map[string]any{service.UpstreamBillingProbeEnabledExtraKey: false}))
  156. stale.Name = "ordinary-edit"
  157. require.NoError(t, repo.Update(ctx, stale))
  158. got, err := repo.GetByID(ctx, account.ID)
  159. require.NoError(t, err)
  160. require.Equal(t, false, got.Extra[service.UpstreamBillingProbeEnabledExtraKey])
  161. require.NotContains(t, got.Extra, service.UpstreamBillingProbeExtraKey)
  162. }
  163. func TestAccountUpdateClearsProbeSnapshotWhenIdentityChanges(t *testing.T) {
  164. ctx := context.Background()
  165. tx := testEntTx(t)
  166. repo := newAccountRepositoryWithSQL(tx.Client(), tx, nil)
  167. account := mustCreateAccount(t, tx.Client(), &service.Account{
  168. Name: "probe-update-identity",
  169. Platform: service.PlatformOpenAI,
  170. Type: service.AccountTypeAPIKey,
  171. Credentials: map[string]any{"api_key": "sk-old"},
  172. Extra: map[string]any{
  173. service.UpstreamBillingProbeEnabledExtraKey: true,
  174. service.UpstreamBillingProbeExtraKey: map[string]any{"status": service.UpstreamBillingProbeStatusOK},
  175. },
  176. })
  177. loaded, err := repo.GetByID(ctx, account.ID)
  178. require.NoError(t, err)
  179. loaded.Credentials["api_key"] = "sk-new"
  180. require.NoError(t, repo.Update(ctx, loaded))
  181. got, err := repo.GetByID(ctx, account.ID)
  182. require.NoError(t, err)
  183. require.NotContains(t, got.Extra, service.UpstreamBillingProbeExtraKey)
  184. }
  185. func TestBulkUpdateAndCredentialUpdateDeleteProbeKey(t *testing.T) {
  186. ctx := context.Background()
  187. tx := testEntTx(t)
  188. repo := newAccountRepositoryWithSQL(tx.Client(), tx, nil)
  189. newAccount := func(name string) *service.Account {
  190. return mustCreateAccount(t, tx.Client(), &service.Account{
  191. Name: name,
  192. Platform: service.PlatformOpenAI,
  193. Type: service.AccountTypeAPIKey,
  194. Credentials: map[string]any{"api_key": "sk-old"},
  195. Extra: map[string]any{
  196. service.UpstreamBillingProbeEnabledExtraKey: true,
  197. service.UpstreamBillingProbeExtraKey: map[string]any{"status": service.UpstreamBillingProbeStatusOK},
  198. },
  199. })
  200. }
  201. bulkAccount := newAccount("probe-bulk-clear")
  202. _, err := repo.BulkUpdate(ctx, []int64{bulkAccount.ID}, service.AccountBulkUpdate{
  203. Extra: map[string]any{service.UpstreamBillingProbeExtraKey: nil},
  204. })
  205. require.NoError(t, err)
  206. got, err := repo.GetByID(ctx, bulkAccount.ID)
  207. require.NoError(t, err)
  208. require.NotContains(t, got.Extra, service.UpstreamBillingProbeExtraKey)
  209. credentialAccount := newAccount("probe-credentials-clear")
  210. require.NoError(t, repo.UpdateCredentials(ctx, credentialAccount.ID, map[string]any{"api_key": "sk-new"}))
  211. got, err = repo.GetByID(ctx, credentialAccount.ID)
  212. require.NoError(t, err)
  213. require.NotContains(t, got.Extra, service.UpstreamBillingProbeExtraKey)
  214. }
  215. func TestProbeSnapshotCASIncludesLoadedEnabledState(t *testing.T) {
  216. tests := []struct {
  217. name string
  218. loadedEnabled bool
  219. concurrentFlip *bool
  220. wantConflict bool
  221. }{
  222. {name: "manual_false_stays_false", loadedEnabled: false},
  223. {name: "periodic_true_disabled_in_flight", loadedEnabled: true, concurrentFlip: boolPtr(false), wantConflict: true},
  224. {name: "manual_false_enabled_in_flight", loadedEnabled: false, concurrentFlip: boolPtr(true), wantConflict: true},
  225. }
  226. for _, tt := range tests {
  227. t.Run(tt.name, func(t *testing.T) {
  228. ctx := context.Background()
  229. tx := testEntTx(t)
  230. repo := newAccountRepositoryWithSQL(tx.Client(), tx, nil)
  231. account := mustCreateAccount(t, tx.Client(), &service.Account{
  232. Name: "probe-enabled-cas-" + tt.name,
  233. Platform: service.PlatformOpenAI,
  234. Type: service.AccountTypeAPIKey,
  235. Credentials: map[string]any{"api_key": "sk-test"},
  236. Extra: map[string]any{service.UpstreamBillingProbeEnabledExtraKey: tt.loadedEnabled},
  237. })
  238. inFlight, err := repo.GetByID(ctx, account.ID)
  239. require.NoError(t, err)
  240. if tt.concurrentFlip != nil {
  241. require.NoError(t, repo.UpdateExtra(ctx, account.ID, map[string]any{
  242. service.UpstreamBillingProbeEnabledExtraKey: *tt.concurrentFlip,
  243. }))
  244. }
  245. err = repo.UpdateUpstreamBillingProbeSnapshot(ctx, inFlight, &service.UpstreamBillingProbeSnapshot{
  246. Status: service.UpstreamBillingProbeStatusOK,
  247. LastAttemptAt: time.Now().UTC(),
  248. }, nil)
  249. if tt.wantConflict {
  250. require.ErrorIs(t, err, service.ErrUpstreamBillingProbeIdentityChanged)
  251. } else {
  252. require.NoError(t, err)
  253. }
  254. got, err := repo.GetByID(ctx, account.ID)
  255. require.NoError(t, err)
  256. if tt.wantConflict {
  257. require.NotContains(t, got.Extra, service.UpstreamBillingProbeExtraKey)
  258. } else {
  259. require.Contains(t, got.Extra, service.UpstreamBillingProbeExtraKey)
  260. }
  261. })
  262. }
  263. }
  264. func TestProbeSnapshotCASProtectsManualRateAfterSyncDisabled(t *testing.T) {
  265. ctx := context.Background()
  266. tx := testEntTx(t)
  267. repo := newAccountRepositoryWithSQL(tx.Client(), tx, nil)
  268. initialRate := 0.25
  269. account := mustCreateAccount(t, tx.Client(), &service.Account{
  270. Name: "probe-sync-cas",
  271. Platform: service.PlatformAnthropic,
  272. Type: service.AccountTypeAPIKey,
  273. RateMultiplier: &initialRate,
  274. Credentials: map[string]any{"api_key": "sk-test"},
  275. Extra: map[string]any{
  276. service.UpstreamBillingProbeEnabledExtraKey: true,
  277. service.UpstreamBillingRateSyncEnabledExtraKey: true,
  278. },
  279. })
  280. inFlight, err := repo.GetByID(ctx, account.ID)
  281. require.NoError(t, err)
  282. manual, err := repo.GetByID(ctx, account.ID)
  283. require.NoError(t, err)
  284. manualRate := 0.8
  285. syncDisabled := false
  286. require.NoError(t, repo.UpdateWithAccountBillingSettings(ctx, manual, nil, &syncDisabled, &manualRate))
  287. probedRate := 0.1
  288. err = repo.UpdateUpstreamBillingProbeSnapshot(ctx, inFlight, &service.UpstreamBillingProbeSnapshot{
  289. Status: service.UpstreamBillingProbeStatusOK,
  290. LastAttemptAt: time.Now().UTC(),
  291. }, &probedRate)
  292. require.ErrorIs(t, err, service.ErrUpstreamBillingProbeIdentityChanged)
  293. got, err := repo.GetByID(ctx, account.ID)
  294. require.NoError(t, err)
  295. require.NotNil(t, got.RateMultiplier)
  296. require.Equal(t, manualRate, *got.RateMultiplier)
  297. require.Equal(t, false, got.Extra[service.UpstreamBillingRateSyncEnabledExtraKey])
  298. require.NotContains(t, got.Extra, service.UpstreamBillingProbeExtraKey)
  299. }
  300. func boolPtr(value bool) *bool {
  301. return &value
  302. }
  303. func TestProxyIdentityUpdateInvalidatesProbeAndRejectsInFlightSnapshot(t *testing.T) {
  304. tests := []struct {
  305. name string
  306. includeProbeKey bool
  307. probeValue any
  308. wantInvalidation bool
  309. }{
  310. {name: "missing_snapshot"},
  311. {name: "json_null_snapshot", includeProbeKey: true},
  312. {name: "existing_snapshot", includeProbeKey: true, probeValue: map[string]any{"status": service.UpstreamBillingProbeStatusOK}, wantInvalidation: true},
  313. }
  314. for _, tt := range tests {
  315. t.Run(tt.name, func(t *testing.T) {
  316. ctx := context.Background()
  317. tx := testEntTx(t)
  318. accountRepo := newAccountRepositoryWithSQL(tx.Client(), tx, nil)
  319. proxyRepo := newProxyRepositoryWithSQL(tx.Client(), tx)
  320. proxy := mustCreateProxy(t, tx.Client(), &service.Proxy{
  321. Name: "probe-proxy",
  322. Protocol: "http",
  323. Host: "old.example",
  324. Port: 8080,
  325. Username: "old-user",
  326. Password: "old-pass",
  327. Status: service.StatusActive,
  328. })
  329. extra := map[string]any{service.UpstreamBillingProbeEnabledExtraKey: true}
  330. if tt.includeProbeKey {
  331. extra[service.UpstreamBillingProbeExtraKey] = tt.probeValue
  332. }
  333. account := mustCreateAccount(t, tx.Client(), &service.Account{
  334. Name: "proxy-probe-account",
  335. Platform: service.PlatformOpenAI,
  336. Type: service.AccountTypeAPIKey,
  337. Credentials: map[string]any{"api_key": "sk-test"},
  338. Extra: extra,
  339. ProxyID: &proxy.ID,
  340. })
  341. inFlight, err := accountRepo.GetByID(ctx, account.ID)
  342. require.NoError(t, err)
  343. require.NotNil(t, inFlight.Proxy)
  344. require.Equal(t, "old.example", inFlight.Proxy.Host)
  345. proxyToUpdate, err := proxyRepo.GetByID(ctx, proxy.ID)
  346. require.NoError(t, err)
  347. proxyToUpdate.Host = "new.example"
  348. require.NoError(t, proxyRepo.Update(ctx, proxyToUpdate))
  349. got, err := accountRepo.GetByID(ctx, account.ID)
  350. require.NoError(t, err)
  351. if tt.wantInvalidation || !tt.includeProbeKey {
  352. require.NotContains(t, got.Extra, service.UpstreamBillingProbeExtraKey)
  353. } else {
  354. require.Contains(t, got.Extra, service.UpstreamBillingProbeExtraKey)
  355. require.Nil(t, got.Extra[service.UpstreamBillingProbeExtraKey])
  356. }
  357. if !tt.wantInvalidation {
  358. require.Equal(t, inFlight.UpdatedAt, got.UpdatedAt, "missing/null snapshots must not cause an account row write")
  359. }
  360. err = accountRepo.UpdateUpstreamBillingProbeSnapshot(ctx, inFlight, &service.UpstreamBillingProbeSnapshot{
  361. Status: service.UpstreamBillingProbeStatusOK,
  362. LastAttemptAt: time.Now().UTC(),
  363. }, nil)
  364. require.ErrorIs(t, err, service.ErrUpstreamBillingProbeIdentityChanged)
  365. rows, err := tx.QueryContext(ctx, `
  366. SELECT COUNT(*), COALESCE(MAX(payload::text), '')
  367. FROM scheduler_outbox
  368. WHERE event_type = $1
  369. `, service.SchedulerOutboxEventAccountBulkChanged)
  370. require.NoError(t, err)
  371. require.True(t, rows.Next())
  372. var (
  373. outboxCount int
  374. payloadJSON string
  375. )
  376. require.NoError(t, rows.Scan(&outboxCount, &payloadJSON))
  377. require.NoError(t, rows.Close())
  378. if tt.wantInvalidation {
  379. require.Equal(t, 1, outboxCount)
  380. var payload struct {
  381. AccountIDs []int64 `json:"account_ids"`
  382. }
  383. require.NoError(t, json.Unmarshal([]byte(payloadJSON), &payload))
  384. require.Equal(t, []int64{account.ID}, payload.AccountIDs)
  385. } else {
  386. require.Zero(t, outboxCount, "no snapshot change means no PR2 cache invalidation event")
  387. }
  388. })
  389. }
  390. }
  391. func TestSweepExpiredProxyWithoutFallbackInvalidatesOnlyExistingProbeSnapshot(t *testing.T) {
  392. ctx := context.Background()
  393. tx := testEntTx(t)
  394. proxyRepo := newProxyRepositoryWithSQL(tx.Client(), tx)
  395. accountRepo := newAccountRepositoryWithSQL(tx.Client(), tx, nil)
  396. past := time.Now().Add(-time.Hour)
  397. proxy := &service.Proxy{
  398. Name: "expired-probe-proxy-none",
  399. Protocol: "http",
  400. Host: "127.0.0.1",
  401. Port: 8080,
  402. Status: service.StatusActive,
  403. ExpiresAt: &past,
  404. FallbackMode: service.FallbackModeNone,
  405. ExpiryWarnDays: 7,
  406. }
  407. require.NoError(t, proxyRepo.Create(ctx, proxy))
  408. newAccount := func(name string, probe any, includeProbe bool) *service.Account {
  409. extra := map[string]any{service.UpstreamBillingProbeEnabledExtraKey: true}
  410. if includeProbe {
  411. extra[service.UpstreamBillingProbeExtraKey] = probe
  412. }
  413. return mustCreateAccount(t, tx.Client(), &service.Account{
  414. Name: name,
  415. Platform: service.PlatformOpenAI,
  416. Type: service.AccountTypeAPIKey,
  417. Credentials: map[string]any{"api_key": "sk-test"},
  418. Extra: extra,
  419. ProxyID: &proxy.ID,
  420. })
  421. }
  422. withSnapshot := newAccount("expired-proxy-with-snapshot", map[string]any{"status": service.UpstreamBillingProbeStatusOK}, true)
  423. withoutSnapshot := newAccount("expired-proxy-without-snapshot", nil, false)
  424. withJSONNull := newAccount("expired-proxy-null-snapshot", nil, true)
  425. untouchedUpdatedAt := make(map[int64]time.Time, 2)
  426. for _, untouched := range []*service.Account{withoutSnapshot, withJSONNull} {
  427. loaded, err := accountRepo.GetByID(ctx, untouched.ID)
  428. require.NoError(t, err)
  429. untouchedUpdatedAt[untouched.ID] = loaded.UpdatedAt
  430. }
  431. changed, err := proxyRepo.SweepExpiredProxies(ctx, time.Now())
  432. require.NoError(t, err)
  433. require.Zero(t, changed, "probe invalidation must not inflate the rerouted account count")
  434. got, err := accountRepo.GetByID(ctx, withSnapshot.ID)
  435. require.NoError(t, err)
  436. require.NotContains(t, got.Extra, service.UpstreamBillingProbeExtraKey)
  437. for _, untouched := range []*service.Account{withoutSnapshot, withJSONNull} {
  438. got, err = accountRepo.GetByID(ctx, untouched.ID)
  439. require.NoError(t, err)
  440. require.Equal(t, untouchedUpdatedAt[untouched.ID], got.UpdatedAt)
  441. }
  442. payload := latestBulkAccountOutboxPayload(t, ctx, tx)
  443. require.Equal(t, []int64{withSnapshot.ID}, payload)
  444. }
  445. func TestSweepExpiredProxyFallbackRerouteDeletesProbeSnapshot(t *testing.T) {
  446. ctx := context.Background()
  447. tx := testEntTx(t)
  448. proxyRepo := newProxyRepositoryWithSQL(tx.Client(), tx)
  449. accountRepo := newAccountRepositoryWithSQL(tx.Client(), tx, nil)
  450. past := time.Now().Add(-time.Hour)
  451. proxy := &service.Proxy{
  452. Name: "expired-probe-proxy-direct",
  453. Protocol: "http",
  454. Host: "127.0.0.1",
  455. Port: 8080,
  456. Status: service.StatusActive,
  457. ExpiresAt: &past,
  458. FallbackMode: service.FallbackModeDirect,
  459. ExpiryWarnDays: 7,
  460. }
  461. require.NoError(t, proxyRepo.Create(ctx, proxy))
  462. account := mustCreateAccount(t, tx.Client(), &service.Account{
  463. Name: "expired-proxy-rerouted-snapshot",
  464. Platform: service.PlatformOpenAI,
  465. Type: service.AccountTypeAPIKey,
  466. Credentials: map[string]any{"api_key": "sk-test"},
  467. Extra: map[string]any{
  468. service.UpstreamBillingProbeEnabledExtraKey: true,
  469. service.UpstreamBillingProbeExtraKey: map[string]any{"status": service.UpstreamBillingProbeStatusOK},
  470. },
  471. ProxyID: &proxy.ID,
  472. })
  473. changed, err := proxyRepo.SweepExpiredProxies(ctx, time.Now())
  474. require.NoError(t, err)
  475. require.EqualValues(t, 1, changed)
  476. got, err := accountRepo.GetByID(ctx, account.ID)
  477. require.NoError(t, err)
  478. require.Nil(t, got.ProxyID)
  479. require.NotContains(t, got.Extra, service.UpstreamBillingProbeExtraKey)
  480. require.Equal(t, []int64{account.ID}, latestBulkAccountOutboxPayload(t, ctx, tx))
  481. }
  482. func latestBulkAccountOutboxPayload(t *testing.T, ctx context.Context, tx sqlQueryer) []int64 {
  483. t.Helper()
  484. var payloadJSON []byte
  485. require.NoError(t, scanSingleRow(ctx, tx, `
  486. SELECT payload
  487. FROM scheduler_outbox
  488. WHERE event_type = $1
  489. ORDER BY id DESC
  490. LIMIT 1
  491. `, []any{service.SchedulerOutboxEventAccountBulkChanged}, &payloadJSON))
  492. var payload struct {
  493. AccountIDs []int64 `json:"account_ids"`
  494. }
  495. require.NoError(t, json.Unmarshal(payloadJSON, &payload))
  496. return payload.AccountIDs
  497. }