usage_log_repo_integration_test.go 63 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091929394959697989910010110210310410510610710810911011111211311411511611711811912012112212312412512612712812913013113213313413513613713813914014114214314414514614714814915015115215315415515615715815916016116216316416516616716816917017117217317417517617717817918018118218318418518618718818919019119219319419519619719819920020120220320420520620720820921021121221321421521621721821922022122222322422522622722822923023123223323423523623723823924024124224324424524624724824925025125225325425525625725825926026126226326426526626726826927027127227327427527627727827928028128228328428528628728828929029129229329429529629729829930030130230330430530630730830931031131231331431531631731831932032132232332432532632732832933033133233333433533633733833934034134234334434534634734834935035135235335435535635735835936036136236336436536636736836937037137237337437537637737837938038138238338438538638738838939039139239339439539639739839940040140240340440540640740840941041141241341441541641741841942042142242342442542642742842943043143243343443543643743843944044144244344444544644744844945045145245345445545645745845946046146246346446546646746846947047147247347447547647747847948048148248348448548648748848949049149249349449549649749849950050150250350450550650750850951051151251351451551651751851952052152252352452552652752852953053153253353453553653753853954054154254354454554654754854955055155255355455555655755855956056156256356456556656756856957057157257357457557657757857958058158258358458558658758858959059159259359459559659759859960060160260360460560660760860961061161261361461561661761861962062162262362462562662762862963063163263363463563663763863964064164264364464564664764864965065165265365465565665765865966066166266366466566666766866967067167267367467567667767867968068168268368468568668768868969069169269369469569669769869970070170270370470570670770870971071171271371471571671771871972072172272372472572672772872973073173273373473573673773873974074174274374474574674774874975075175275375475575675775875976076176276376476576676776876977077177277377477577677777877978078178278378478578678778878979079179279379479579679779879980080180280380480580680780880981081181281381481581681781881982082182282382482582682782882983083183283383483583683783883984084184284384484584684784884985085185285385485585685785885986086186286386486586686786886987087187287387487587687787887988088188288388488588688788888989089189289389489589689789889990090190290390490590690790890991091191291391491591691791891992092192292392492592692792892993093193293393493593693793893994094194294394494594694794894995095195295395495595695795895996096196296396496596696796896997097197297397497597697797897998098198298398498598698798898999099199299399499599699799899910001001100210031004100510061007100810091010101110121013101410151016101710181019102010211022102310241025102610271028102910301031103210331034103510361037103810391040104110421043104410451046104710481049105010511052105310541055105610571058105910601061106210631064106510661067106810691070107110721073107410751076107710781079108010811082108310841085108610871088108910901091109210931094109510961097109810991100110111021103110411051106110711081109111011111112111311141115111611171118111911201121112211231124112511261127112811291130113111321133113411351136113711381139114011411142114311441145114611471148114911501151115211531154115511561157115811591160116111621163116411651166116711681169117011711172117311741175117611771178117911801181118211831184118511861187118811891190119111921193119411951196119711981199120012011202120312041205120612071208120912101211121212131214121512161217121812191220122112221223122412251226122712281229123012311232123312341235123612371238123912401241124212431244124512461247124812491250125112521253125412551256125712581259126012611262126312641265126612671268126912701271127212731274127512761277127812791280128112821283128412851286128712881289129012911292129312941295129612971298129913001301130213031304130513061307130813091310131113121313131413151316131713181319132013211322132313241325132613271328132913301331133213331334133513361337133813391340134113421343134413451346134713481349135013511352135313541355135613571358135913601361136213631364136513661367136813691370137113721373137413751376137713781379138013811382138313841385138613871388138913901391139213931394139513961397139813991400140114021403140414051406140714081409141014111412141314141415141614171418141914201421142214231424142514261427142814291430143114321433143414351436143714381439144014411442144314441445144614471448144914501451145214531454145514561457145814591460146114621463146414651466146714681469147014711472147314741475147614771478147914801481148214831484148514861487148814891490149114921493149414951496149714981499150015011502150315041505150615071508150915101511151215131514151515161517151815191520152115221523152415251526152715281529153015311532153315341535153615371538153915401541154215431544154515461547154815491550155115521553155415551556155715581559156015611562156315641565156615671568156915701571157215731574157515761577157815791580158115821583158415851586158715881589159015911592159315941595159615971598159916001601160216031604160516061607160816091610161116121613161416151616161716181619162016211622162316241625162616271628162916301631163216331634163516361637163816391640164116421643164416451646164716481649165016511652165316541655165616571658165916601661166216631664166516661667166816691670167116721673167416751676167716781679
  1. //go:build integration
  2. package repository
  3. import (
  4. "context"
  5. "fmt"
  6. "sync"
  7. "testing"
  8. "time"
  9. "github.com/google/uuid"
  10. dbent "github.com/Wei-Shaw/sub2api/ent"
  11. "github.com/Wei-Shaw/sub2api/internal/pkg/pagination"
  12. "github.com/Wei-Shaw/sub2api/internal/pkg/timezone"
  13. "github.com/Wei-Shaw/sub2api/internal/pkg/usagestats"
  14. "github.com/Wei-Shaw/sub2api/internal/service"
  15. "github.com/stretchr/testify/require"
  16. "github.com/stretchr/testify/suite"
  17. )
  18. type UsageLogRepoSuite struct {
  19. suite.Suite
  20. ctx context.Context
  21. tx *dbent.Tx
  22. client *dbent.Client
  23. repo *usageLogRepository
  24. }
  25. func (s *UsageLogRepoSuite) SetupTest() {
  26. s.ctx = context.Background()
  27. tx := testEntTx(s.T())
  28. s.tx = tx
  29. s.client = tx.Client()
  30. s.repo = newUsageLogRepositoryWithSQL(s.client, tx)
  31. }
  32. func TestUsageLogRepoSuite(t *testing.T) {
  33. suite.Run(t, new(UsageLogRepoSuite))
  34. }
  35. // truncateToDayUTC 截断到 UTC 日期边界(测试辅助函数)
  36. func truncateToDayUTC(t time.Time) time.Time {
  37. t = t.UTC()
  38. return time.Date(t.Year(), t.Month(), t.Day(), 0, 0, 0, 0, time.UTC)
  39. }
  40. func (s *UsageLogRepoSuite) createUsageLog(user *service.User, apiKey *service.APIKey, account *service.Account, inputTokens, outputTokens int, cost float64, createdAt time.Time) *service.UsageLog {
  41. log := &service.UsageLog{
  42. UserID: user.ID,
  43. APIKeyID: apiKey.ID,
  44. AccountID: account.ID,
  45. RequestID: uuid.New().String(), // Generate unique RequestID for each log
  46. Model: "claude-3",
  47. InputTokens: inputTokens,
  48. OutputTokens: outputTokens,
  49. TotalCost: cost,
  50. ActualCost: cost,
  51. CreatedAt: createdAt,
  52. }
  53. _, err := s.repo.Create(s.ctx, log)
  54. s.Require().NoError(err)
  55. return log
  56. }
  57. // --- Create / GetByID ---
  58. func (s *UsageLogRepoSuite) TestCreate() {
  59. user := mustCreateUser(s.T(), s.client, &service.User{Email: "create@test.com"})
  60. apiKey := mustCreateApiKey(s.T(), s.client, &service.APIKey{UserID: user.ID, Key: "sk-create", Name: "k"})
  61. account := mustCreateAccount(s.T(), s.client, &service.Account{Name: "acc-create"})
  62. log := &service.UsageLog{
  63. UserID: user.ID,
  64. APIKeyID: apiKey.ID,
  65. AccountID: account.ID,
  66. Model: "claude-3",
  67. InputTokens: 10,
  68. OutputTokens: 20,
  69. TotalCost: 0.5,
  70. ActualCost: 0.4,
  71. }
  72. _, err := s.repo.Create(s.ctx, log)
  73. s.Require().NoError(err, "Create")
  74. s.Require().NotZero(log.ID)
  75. }
  76. func TestUsageLogRepositoryCreate_BatchPathConcurrent(t *testing.T) {
  77. ctx := context.Background()
  78. client := testEntClient(t)
  79. repo := newUsageLogRepositoryWithSQL(client, integrationDB)
  80. user := mustCreateUser(t, client, &service.User{Email: fmt.Sprintf("usage-batch-%d@example.com", time.Now().UnixNano())})
  81. apiKey := mustCreateApiKey(t, client, &service.APIKey{UserID: user.ID, Key: "sk-usage-batch-" + uuid.NewString(), Name: "k"})
  82. account := mustCreateAccount(t, client, &service.Account{Name: "acc-usage-batch-" + uuid.NewString()})
  83. const total = 16
  84. results := make([]bool, total)
  85. errs := make([]error, total)
  86. logs := make([]*service.UsageLog, total)
  87. var wg sync.WaitGroup
  88. wg.Add(total)
  89. for i := 0; i < total; i++ {
  90. i := i
  91. logs[i] = &service.UsageLog{
  92. UserID: user.ID,
  93. APIKeyID: apiKey.ID,
  94. AccountID: account.ID,
  95. RequestID: uuid.NewString(),
  96. Model: "claude-3",
  97. InputTokens: 10 + i,
  98. OutputTokens: 20 + i,
  99. TotalCost: 0.5,
  100. ActualCost: 0.5,
  101. CreatedAt: time.Now().UTC(),
  102. }
  103. go func() {
  104. defer wg.Done()
  105. results[i], errs[i] = repo.Create(ctx, logs[i])
  106. }()
  107. }
  108. wg.Wait()
  109. for i := 0; i < total; i++ {
  110. require.NoError(t, errs[i])
  111. require.True(t, results[i])
  112. require.NotZero(t, logs[i].ID)
  113. }
  114. var count int
  115. require.NoError(t, integrationDB.QueryRowContext(ctx, "SELECT COUNT(*) FROM usage_logs WHERE api_key_id = $1", apiKey.ID).Scan(&count))
  116. require.Equal(t, total, count)
  117. }
  118. func TestUsageLogRepositoryCreate_BatchPathDuplicateRequestID(t *testing.T) {
  119. ctx := context.Background()
  120. client := testEntClient(t)
  121. repo := newUsageLogRepositoryWithSQL(client, integrationDB)
  122. user := mustCreateUser(t, client, &service.User{Email: fmt.Sprintf("usage-dup-%d@example.com", time.Now().UnixNano())})
  123. apiKey := mustCreateApiKey(t, client, &service.APIKey{UserID: user.ID, Key: "sk-usage-dup-" + uuid.NewString(), Name: "k"})
  124. account := mustCreateAccount(t, client, &service.Account{Name: "acc-usage-dup-" + uuid.NewString()})
  125. requestID := uuid.NewString()
  126. log1 := &service.UsageLog{
  127. UserID: user.ID,
  128. APIKeyID: apiKey.ID,
  129. AccountID: account.ID,
  130. RequestID: requestID,
  131. Model: "claude-3",
  132. InputTokens: 10,
  133. OutputTokens: 20,
  134. TotalCost: 0.5,
  135. ActualCost: 0.5,
  136. CreatedAt: time.Now().UTC(),
  137. }
  138. log2 := &service.UsageLog{
  139. UserID: user.ID,
  140. APIKeyID: apiKey.ID,
  141. AccountID: account.ID,
  142. RequestID: requestID,
  143. Model: "claude-3",
  144. InputTokens: 10,
  145. OutputTokens: 20,
  146. TotalCost: 0.5,
  147. ActualCost: 0.5,
  148. CreatedAt: time.Now().UTC(),
  149. }
  150. inserted1, err1 := repo.Create(ctx, log1)
  151. inserted2, err2 := repo.Create(ctx, log2)
  152. require.NoError(t, err1)
  153. require.NoError(t, err2)
  154. require.True(t, inserted1)
  155. require.False(t, inserted2)
  156. require.Equal(t, log1.ID, log2.ID)
  157. var count int
  158. require.NoError(t, integrationDB.QueryRowContext(ctx, "SELECT COUNT(*) FROM usage_logs WHERE request_id = $1 AND api_key_id = $2", requestID, apiKey.ID).Scan(&count))
  159. require.Equal(t, 1, count)
  160. }
  161. func TestUsageLogRepositoryFlushCreateBatch_DeduplicatesSameKeyInMemory(t *testing.T) {
  162. ctx := context.Background()
  163. client := testEntClient(t)
  164. repo := newUsageLogRepositoryWithSQL(client, integrationDB)
  165. user := mustCreateUser(t, client, &service.User{Email: fmt.Sprintf("usage-batch-memdup-%d@example.com", time.Now().UnixNano())})
  166. apiKey := mustCreateApiKey(t, client, &service.APIKey{UserID: user.ID, Key: "sk-usage-batch-memdup-" + uuid.NewString(), Name: "k"})
  167. account := mustCreateAccount(t, client, &service.Account{Name: "acc-usage-batch-memdup-" + uuid.NewString()})
  168. requestID := uuid.NewString()
  169. const total = 8
  170. batch := make([]usageLogCreateRequest, 0, total)
  171. logs := make([]*service.UsageLog, 0, total)
  172. for i := 0; i < total; i++ {
  173. log := &service.UsageLog{
  174. UserID: user.ID,
  175. APIKeyID: apiKey.ID,
  176. AccountID: account.ID,
  177. RequestID: requestID,
  178. Model: "claude-3",
  179. InputTokens: 10 + i,
  180. OutputTokens: 20 + i,
  181. TotalCost: 0.5,
  182. ActualCost: 0.5,
  183. CreatedAt: time.Now().UTC(),
  184. }
  185. logs = append(logs, log)
  186. batch = append(batch, usageLogCreateRequest{
  187. log: log,
  188. prepared: prepareUsageLogInsert(log),
  189. resultCh: make(chan usageLogCreateResult, 1),
  190. })
  191. }
  192. repo.flushCreateBatch(integrationDB, batch)
  193. insertedCount := 0
  194. var firstID int64
  195. for idx, req := range batch {
  196. res := <-req.resultCh
  197. require.NoError(t, res.err)
  198. if res.inserted {
  199. insertedCount++
  200. }
  201. require.NotZero(t, logs[idx].ID)
  202. if idx == 0 {
  203. firstID = logs[idx].ID
  204. } else {
  205. require.Equal(t, firstID, logs[idx].ID)
  206. }
  207. }
  208. require.Equal(t, 1, insertedCount)
  209. var count int
  210. require.NoError(t, integrationDB.QueryRowContext(ctx, "SELECT COUNT(*) FROM usage_logs WHERE request_id = $1 AND api_key_id = $2", requestID, apiKey.ID).Scan(&count))
  211. require.Equal(t, 1, count)
  212. }
  213. func TestUsageLogRepositoryCreateBestEffort_BatchPathDuplicateRequestID(t *testing.T) {
  214. ctx := context.Background()
  215. client := testEntClient(t)
  216. repo := newUsageLogRepositoryWithSQL(client, integrationDB)
  217. user := mustCreateUser(t, client, &service.User{Email: fmt.Sprintf("usage-best-effort-dup-%d@example.com", time.Now().UnixNano())})
  218. apiKey := mustCreateApiKey(t, client, &service.APIKey{UserID: user.ID, Key: "sk-usage-best-effort-dup-" + uuid.NewString(), Name: "k"})
  219. account := mustCreateAccount(t, client, &service.Account{Name: "acc-usage-best-effort-dup-" + uuid.NewString()})
  220. requestID := uuid.NewString()
  221. log1 := &service.UsageLog{
  222. UserID: user.ID,
  223. APIKeyID: apiKey.ID,
  224. AccountID: account.ID,
  225. RequestID: requestID,
  226. Model: "claude-3",
  227. InputTokens: 10,
  228. OutputTokens: 20,
  229. TotalCost: 0.5,
  230. ActualCost: 0.5,
  231. CreatedAt: time.Now().UTC(),
  232. }
  233. log2 := &service.UsageLog{
  234. UserID: user.ID,
  235. APIKeyID: apiKey.ID,
  236. AccountID: account.ID,
  237. RequestID: requestID,
  238. Model: "claude-3",
  239. InputTokens: 10,
  240. OutputTokens: 20,
  241. TotalCost: 0.5,
  242. ActualCost: 0.5,
  243. CreatedAt: time.Now().UTC(),
  244. }
  245. require.NoError(t, repo.CreateBestEffort(ctx, log1))
  246. require.NoError(t, repo.CreateBestEffort(ctx, log2))
  247. require.Eventually(t, func() bool {
  248. var count int
  249. err := integrationDB.QueryRowContext(ctx, "SELECT COUNT(*) FROM usage_logs WHERE request_id = $1 AND api_key_id = $2", requestID, apiKey.ID).Scan(&count)
  250. return err == nil && count == 1
  251. }, 3*time.Second, 20*time.Millisecond)
  252. }
  253. func TestUsageLogRepositoryCreateBestEffort_QueueFullBlocksUntilCtxDeadline(t *testing.T) {
  254. // 队列满时不再立即丢弃:阻塞等待入队,直到调用方 ctx 到期才标记 dropped(issue #3656)。
  255. client := testEntClient(t)
  256. repo := newUsageLogRepositoryWithSQL(client, integrationDB)
  257. repo.bestEffortBatchCh = make(chan usageLogBestEffortRequest, 1)
  258. repo.bestEffortBatchCh <- usageLogBestEffortRequest{}
  259. ctx, cancel := context.WithTimeout(context.Background(), 200*time.Millisecond)
  260. defer cancel()
  261. start := time.Now()
  262. err := repo.CreateBestEffort(ctx, &service.UsageLog{
  263. UserID: 1,
  264. APIKeyID: 2,
  265. AccountID: 3,
  266. RequestID: uuid.NewString(),
  267. Model: "claude-3",
  268. InputTokens: 10,
  269. OutputTokens: 20,
  270. TotalCost: 0.5,
  271. ActualCost: 0.5,
  272. CreatedAt: time.Now().UTC(),
  273. })
  274. require.Error(t, err)
  275. require.True(t, service.IsUsageLogCreateDropped(err))
  276. require.GreaterOrEqual(t, time.Since(start), 150*time.Millisecond)
  277. }
  278. func TestUsageLogRepositoryCreateBestEffort_QueueFullWaitsForDrain(t *testing.T) {
  279. // 队列满但批处理器随后排空时,阻塞的入队应成功完成而非丢弃。
  280. client := testEntClient(t)
  281. repo := newUsageLogRepositoryWithSQL(client, integrationDB)
  282. repo.bestEffortBatchCh = make(chan usageLogBestEffortRequest, 1)
  283. repo.bestEffortBatchCh <- usageLogBestEffortRequest{}
  284. go func() {
  285. time.Sleep(100 * time.Millisecond)
  286. <-repo.bestEffortBatchCh // 排空占位请求,为阻塞中的入队腾出空间
  287. req := <-repo.bestEffortBatchCh
  288. sendUsageLogBestEffortResult(req.resultCh, nil)
  289. }()
  290. ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
  291. defer cancel()
  292. err := repo.CreateBestEffort(ctx, &service.UsageLog{
  293. UserID: 1,
  294. APIKeyID: 2,
  295. AccountID: 3,
  296. RequestID: uuid.NewString(),
  297. Model: "claude-3",
  298. InputTokens: 10,
  299. OutputTokens: 20,
  300. TotalCost: 0.5,
  301. ActualCost: 0.5,
  302. CreatedAt: time.Now().UTC(),
  303. })
  304. require.NoError(t, err)
  305. }
  306. func TestUsageLogRepositoryCreate_BatchPathCanceledContextMarksNotPersisted(t *testing.T) {
  307. client := testEntClient(t)
  308. repo := newUsageLogRepositoryWithSQL(client, integrationDB)
  309. user := mustCreateUser(t, client, &service.User{Email: fmt.Sprintf("usage-cancel-%d@example.com", time.Now().UnixNano())})
  310. apiKey := mustCreateApiKey(t, client, &service.APIKey{UserID: user.ID, Key: "sk-usage-cancel-" + uuid.NewString(), Name: "k"})
  311. account := mustCreateAccount(t, client, &service.Account{Name: "acc-usage-cancel-" + uuid.NewString()})
  312. ctx, cancel := context.WithCancel(context.Background())
  313. cancel()
  314. inserted, err := repo.Create(ctx, &service.UsageLog{
  315. UserID: user.ID,
  316. APIKeyID: apiKey.ID,
  317. AccountID: account.ID,
  318. RequestID: uuid.NewString(),
  319. Model: "claude-3",
  320. InputTokens: 10,
  321. OutputTokens: 20,
  322. TotalCost: 0.5,
  323. ActualCost: 0.5,
  324. CreatedAt: time.Now().UTC(),
  325. })
  326. require.False(t, inserted)
  327. require.Error(t, err)
  328. require.True(t, service.IsUsageLogCreateNotPersisted(err))
  329. }
  330. func TestUsageLogRepositoryCreate_BatchPathQueueFullMarksNotPersisted(t *testing.T) {
  331. // 队列满时阻塞等待入队,直到调用方 ctx 到期才标记 not persisted(issue #3656)。
  332. client := testEntClient(t)
  333. repo := newUsageLogRepositoryWithSQL(client, integrationDB)
  334. repo.createBatchCh = make(chan usageLogCreateRequest, 1)
  335. repo.createBatchCh <- usageLogCreateRequest{}
  336. user := mustCreateUser(t, client, &service.User{Email: fmt.Sprintf("usage-create-full-%d@example.com", time.Now().UnixNano())})
  337. apiKey := mustCreateApiKey(t, client, &service.APIKey{UserID: user.ID, Key: "sk-usage-create-full-" + uuid.NewString(), Name: "k"})
  338. account := mustCreateAccount(t, client, &service.Account{Name: "acc-usage-create-full-" + uuid.NewString()})
  339. ctx, cancel := context.WithTimeout(context.Background(), 200*time.Millisecond)
  340. defer cancel()
  341. start := time.Now()
  342. inserted, err := repo.Create(ctx, &service.UsageLog{
  343. UserID: user.ID,
  344. APIKeyID: apiKey.ID,
  345. AccountID: account.ID,
  346. RequestID: uuid.NewString(),
  347. Model: "claude-3",
  348. InputTokens: 10,
  349. OutputTokens: 20,
  350. TotalCost: 0.5,
  351. ActualCost: 0.5,
  352. CreatedAt: time.Now().UTC(),
  353. })
  354. require.False(t, inserted)
  355. require.Error(t, err)
  356. require.True(t, service.IsUsageLogCreateNotPersisted(err))
  357. require.GreaterOrEqual(t, time.Since(start), 150*time.Millisecond)
  358. }
  359. func TestUsageLogRepositoryCreate_BatchPathCanceledAfterQueueMarksNotPersisted(t *testing.T) {
  360. client := testEntClient(t)
  361. repo := newUsageLogRepositoryWithSQL(client, integrationDB)
  362. repo.createBatchCh = make(chan usageLogCreateRequest, 1)
  363. user := mustCreateUser(t, client, &service.User{Email: fmt.Sprintf("usage-cancel-queued-%d@example.com", time.Now().UnixNano())})
  364. apiKey := mustCreateApiKey(t, client, &service.APIKey{UserID: user.ID, Key: "sk-usage-cancel-queued-" + uuid.NewString(), Name: "k"})
  365. account := mustCreateAccount(t, client, &service.Account{Name: "acc-usage-cancel-queued-" + uuid.NewString()})
  366. ctx, cancel := context.WithCancel(context.Background())
  367. errCh := make(chan error, 1)
  368. go func() {
  369. _, err := repo.createBatched(ctx, &service.UsageLog{
  370. UserID: user.ID,
  371. APIKeyID: apiKey.ID,
  372. AccountID: account.ID,
  373. RequestID: uuid.NewString(),
  374. Model: "claude-3",
  375. InputTokens: 10,
  376. OutputTokens: 20,
  377. TotalCost: 0.5,
  378. ActualCost: 0.5,
  379. CreatedAt: time.Now().UTC(),
  380. })
  381. errCh <- err
  382. }()
  383. req := <-repo.createBatchCh
  384. require.NotNil(t, req.shared)
  385. cancel()
  386. err := <-errCh
  387. require.Error(t, err)
  388. require.True(t, service.IsUsageLogCreateNotPersisted(err))
  389. completeUsageLogCreateRequest(req, usageLogCreateResult{inserted: false, err: service.MarkUsageLogCreateNotPersisted(context.Canceled)})
  390. }
  391. func TestUsageLogRepositoryFlushCreateBatch_CanceledRequestReturnsNotPersisted(t *testing.T) {
  392. client := testEntClient(t)
  393. repo := newUsageLogRepositoryWithSQL(client, integrationDB)
  394. user := mustCreateUser(t, client, &service.User{Email: fmt.Sprintf("usage-flush-cancel-%d@example.com", time.Now().UnixNano())})
  395. apiKey := mustCreateApiKey(t, client, &service.APIKey{UserID: user.ID, Key: "sk-usage-flush-cancel-" + uuid.NewString(), Name: "k"})
  396. account := mustCreateAccount(t, client, &service.Account{Name: "acc-usage-flush-cancel-" + uuid.NewString()})
  397. log := &service.UsageLog{
  398. UserID: user.ID,
  399. APIKeyID: apiKey.ID,
  400. AccountID: account.ID,
  401. RequestID: uuid.NewString(),
  402. Model: "claude-3",
  403. InputTokens: 10,
  404. OutputTokens: 20,
  405. TotalCost: 0.5,
  406. ActualCost: 0.5,
  407. CreatedAt: time.Now().UTC(),
  408. }
  409. req := usageLogCreateRequest{
  410. log: log,
  411. prepared: prepareUsageLogInsert(log),
  412. shared: &usageLogCreateShared{},
  413. resultCh: make(chan usageLogCreateResult, 1),
  414. }
  415. req.shared.state.Store(usageLogCreateStateCanceled)
  416. repo.flushCreateBatch(integrationDB, []usageLogCreateRequest{req})
  417. res := <-req.resultCh
  418. require.False(t, res.inserted)
  419. require.Error(t, res.err)
  420. require.True(t, service.IsUsageLogCreateNotPersisted(res.err))
  421. }
  422. func (s *UsageLogRepoSuite) TestGetByID() {
  423. user := mustCreateUser(s.T(), s.client, &service.User{Email: "getbyid@test.com"})
  424. apiKey := mustCreateApiKey(s.T(), s.client, &service.APIKey{UserID: user.ID, Key: "sk-getbyid", Name: "k"})
  425. account := mustCreateAccount(s.T(), s.client, &service.Account{Name: "acc-getbyid"})
  426. log := s.createUsageLog(user, apiKey, account, 10, 20, 0.5, time.Now())
  427. got, err := s.repo.GetByID(s.ctx, log.ID)
  428. s.Require().NoError(err, "GetByID")
  429. s.Require().Equal(log.ID, got.ID)
  430. s.Require().Equal(10, got.InputTokens)
  431. }
  432. func (s *UsageLogRepoSuite) TestGetByID_NotFound() {
  433. _, err := s.repo.GetByID(s.ctx, 999999)
  434. s.Require().Error(err, "expected error for non-existent ID")
  435. }
  436. func (s *UsageLogRepoSuite) TestGetByID_ReturnsAccountRateMultiplier() {
  437. user := mustCreateUser(s.T(), s.client, &service.User{Email: "getbyid-mult@test.com"})
  438. apiKey := mustCreateApiKey(s.T(), s.client, &service.APIKey{UserID: user.ID, Key: "sk-getbyid-mult", Name: "k"})
  439. account := mustCreateAccount(s.T(), s.client, &service.Account{Name: "acc-getbyid-mult"})
  440. m := 0.5
  441. log := &service.UsageLog{
  442. UserID: user.ID,
  443. APIKeyID: apiKey.ID,
  444. AccountID: account.ID,
  445. RequestID: uuid.New().String(),
  446. Model: "claude-3",
  447. InputTokens: 10,
  448. OutputTokens: 20,
  449. TotalCost: 1.0,
  450. ActualCost: 2.0,
  451. AccountRateMultiplier: &m,
  452. CreatedAt: timezone.Today().Add(2 * time.Hour),
  453. }
  454. _, err := s.repo.Create(s.ctx, log)
  455. s.Require().NoError(err)
  456. got, err := s.repo.GetByID(s.ctx, log.ID)
  457. s.Require().NoError(err)
  458. s.Require().NotNil(got.AccountRateMultiplier)
  459. s.Require().InEpsilon(0.5, *got.AccountRateMultiplier, 0.0001)
  460. }
  461. func (s *UsageLogRepoSuite) TestGetByID_ReturnsOpenAIWSMode() {
  462. user := mustCreateUser(s.T(), s.client, &service.User{Email: "getbyid-ws@test.com"})
  463. apiKey := mustCreateApiKey(s.T(), s.client, &service.APIKey{UserID: user.ID, Key: "sk-getbyid-ws", Name: "k"})
  464. account := mustCreateAccount(s.T(), s.client, &service.Account{Name: "acc-getbyid-ws"})
  465. log := &service.UsageLog{
  466. UserID: user.ID,
  467. APIKeyID: apiKey.ID,
  468. AccountID: account.ID,
  469. RequestID: uuid.New().String(),
  470. Model: "gpt-5.3-codex",
  471. InputTokens: 10,
  472. OutputTokens: 20,
  473. TotalCost: 1.0,
  474. ActualCost: 1.0,
  475. OpenAIWSMode: true,
  476. CreatedAt: timezone.Today().Add(3 * time.Hour),
  477. }
  478. _, err := s.repo.Create(s.ctx, log)
  479. s.Require().NoError(err)
  480. got, err := s.repo.GetByID(s.ctx, log.ID)
  481. s.Require().NoError(err)
  482. s.Require().True(got.OpenAIWSMode)
  483. }
  484. func (s *UsageLogRepoSuite) TestGetByID_ReturnsRequestTypeAndLegacyFallback() {
  485. user := mustCreateUser(s.T(), s.client, &service.User{Email: "getbyid-request-type@test.com"})
  486. apiKey := mustCreateApiKey(s.T(), s.client, &service.APIKey{UserID: user.ID, Key: "sk-getbyid-request-type", Name: "k"})
  487. account := mustCreateAccount(s.T(), s.client, &service.Account{Name: "acc-getbyid-request-type"})
  488. log := &service.UsageLog{
  489. UserID: user.ID,
  490. APIKeyID: apiKey.ID,
  491. AccountID: account.ID,
  492. RequestID: uuid.New().String(),
  493. Model: "gpt-5.3-codex",
  494. RequestType: service.RequestTypeWSV2,
  495. Stream: true,
  496. OpenAIWSMode: false,
  497. InputTokens: 10,
  498. OutputTokens: 20,
  499. TotalCost: 1.0,
  500. ActualCost: 1.0,
  501. CreatedAt: timezone.Today().Add(4 * time.Hour),
  502. }
  503. _, err := s.repo.Create(s.ctx, log)
  504. s.Require().NoError(err)
  505. got, err := s.repo.GetByID(s.ctx, log.ID)
  506. s.Require().NoError(err)
  507. s.Require().Equal(service.RequestTypeWSV2, got.RequestType)
  508. s.Require().True(got.Stream)
  509. s.Require().True(got.OpenAIWSMode)
  510. }
  511. // --- Delete ---
  512. func (s *UsageLogRepoSuite) TestDelete() {
  513. user := mustCreateUser(s.T(), s.client, &service.User{Email: "delete@test.com"})
  514. apiKey := mustCreateApiKey(s.T(), s.client, &service.APIKey{UserID: user.ID, Key: "sk-delete", Name: "k"})
  515. account := mustCreateAccount(s.T(), s.client, &service.Account{Name: "acc-delete"})
  516. log := s.createUsageLog(user, apiKey, account, 10, 20, 0.5, time.Now())
  517. err := s.repo.Delete(s.ctx, log.ID)
  518. s.Require().NoError(err, "Delete")
  519. _, err = s.repo.GetByID(s.ctx, log.ID)
  520. s.Require().Error(err, "expected error after delete")
  521. }
  522. // --- ListByUser ---
  523. func (s *UsageLogRepoSuite) TestListByUser() {
  524. user := mustCreateUser(s.T(), s.client, &service.User{Email: "listbyuser@test.com"})
  525. apiKey := mustCreateApiKey(s.T(), s.client, &service.APIKey{UserID: user.ID, Key: "sk-listbyuser", Name: "k"})
  526. account := mustCreateAccount(s.T(), s.client, &service.Account{Name: "acc-listbyuser"})
  527. s.createUsageLog(user, apiKey, account, 10, 20, 0.5, time.Now())
  528. s.createUsageLog(user, apiKey, account, 15, 25, 0.6, time.Now())
  529. logs, page, err := s.repo.ListByUser(s.ctx, user.ID, pagination.PaginationParams{Page: 1, PageSize: 10})
  530. s.Require().NoError(err, "ListByUser")
  531. s.Require().Len(logs, 2)
  532. s.Require().Equal(int64(2), page.Total)
  533. }
  534. // --- ListByAPIKey ---
  535. func (s *UsageLogRepoSuite) TestListByAPIKey() {
  536. user := mustCreateUser(s.T(), s.client, &service.User{Email: "listbyapikey@test.com"})
  537. apiKey := mustCreateApiKey(s.T(), s.client, &service.APIKey{UserID: user.ID, Key: "sk-listbyapikey", Name: "k"})
  538. account := mustCreateAccount(s.T(), s.client, &service.Account{Name: "acc-listbyapikey"})
  539. s.createUsageLog(user, apiKey, account, 10, 20, 0.5, time.Now())
  540. s.createUsageLog(user, apiKey, account, 15, 25, 0.6, time.Now())
  541. logs, page, err := s.repo.ListByAPIKey(s.ctx, apiKey.ID, pagination.PaginationParams{Page: 1, PageSize: 10})
  542. s.Require().NoError(err, "ListByAPIKey")
  543. s.Require().Len(logs, 2)
  544. s.Require().Equal(int64(2), page.Total)
  545. }
  546. // --- ListByAccount ---
  547. func (s *UsageLogRepoSuite) TestListByAccount() {
  548. user := mustCreateUser(s.T(), s.client, &service.User{Email: "listbyaccount@test.com"})
  549. apiKey := mustCreateApiKey(s.T(), s.client, &service.APIKey{UserID: user.ID, Key: "sk-listbyaccount", Name: "k"})
  550. account := mustCreateAccount(s.T(), s.client, &service.Account{Name: "acc-listbyaccount"})
  551. s.createUsageLog(user, apiKey, account, 10, 20, 0.5, time.Now())
  552. logs, page, err := s.repo.ListByAccount(s.ctx, account.ID, pagination.PaginationParams{Page: 1, PageSize: 10})
  553. s.Require().NoError(err, "ListByAccount")
  554. s.Require().Len(logs, 1)
  555. s.Require().Equal(int64(1), page.Total)
  556. }
  557. // --- GetUserStats ---
  558. func (s *UsageLogRepoSuite) TestGetUserStats() {
  559. user := mustCreateUser(s.T(), s.client, &service.User{Email: "userstats@test.com"})
  560. apiKey := mustCreateApiKey(s.T(), s.client, &service.APIKey{UserID: user.ID, Key: "sk-userstats", Name: "k"})
  561. account := mustCreateAccount(s.T(), s.client, &service.Account{Name: "acc-userstats"})
  562. base := time.Date(2025, 1, 15, 12, 0, 0, 0, time.UTC)
  563. s.createUsageLog(user, apiKey, account, 10, 20, 0.5, base)
  564. s.createUsageLog(user, apiKey, account, 15, 25, 0.6, base.Add(1*time.Hour))
  565. startTime := base.Add(-1 * time.Hour)
  566. endTime := base.Add(2 * time.Hour)
  567. stats, err := s.repo.GetUserStats(s.ctx, user.ID, startTime, endTime)
  568. s.Require().NoError(err, "GetUserStats")
  569. s.Require().Equal(int64(2), stats.TotalRequests)
  570. s.Require().Equal(int64(25), stats.InputTokens)
  571. s.Require().Equal(int64(45), stats.OutputTokens)
  572. }
  573. // --- ListWithFilters ---
  574. func (s *UsageLogRepoSuite) TestListWithFilters() {
  575. user := mustCreateUser(s.T(), s.client, &service.User{Email: "filters@test.com"})
  576. apiKey := mustCreateApiKey(s.T(), s.client, &service.APIKey{UserID: user.ID, Key: "sk-filters", Name: "k"})
  577. account := mustCreateAccount(s.T(), s.client, &service.Account{Name: "acc-filters"})
  578. s.createUsageLog(user, apiKey, account, 10, 20, 0.5, time.Now())
  579. filters := usagestats.UsageLogFilters{UserID: user.ID}
  580. logs, page, err := s.repo.ListWithFilters(s.ctx, pagination.PaginationParams{Page: 1, PageSize: 10}, filters)
  581. s.Require().NoError(err, "ListWithFilters")
  582. s.Require().Len(logs, 1)
  583. s.Require().Equal(int64(1), page.Total)
  584. }
  585. // --- GetDashboardStats ---
  586. func (s *UsageLogRepoSuite) TestDashboardStats_TodayTotalsAndPerformance() {
  587. now := time.Now().UTC()
  588. todayStart := truncateToDayUTC(now)
  589. baseStats, err := s.repo.GetDashboardStats(s.ctx)
  590. s.Require().NoError(err, "GetDashboardStats base")
  591. userToday := mustCreateUser(s.T(), s.client, &service.User{
  592. Email: "today@example.com",
  593. CreatedAt: testMaxTime(todayStart.Add(10*time.Second), now.Add(-10*time.Second)),
  594. UpdatedAt: now,
  595. })
  596. userOld := mustCreateUser(s.T(), s.client, &service.User{
  597. Email: "old@example.com",
  598. CreatedAt: todayStart.Add(-24 * time.Hour),
  599. UpdatedAt: todayStart.Add(-24 * time.Hour),
  600. })
  601. group := mustCreateGroup(s.T(), s.client, &service.Group{Name: "g-ul"})
  602. apiKey1 := mustCreateApiKey(s.T(), s.client, &service.APIKey{UserID: userToday.ID, Key: "sk-ul-1", Name: "ul1"})
  603. mustCreateApiKey(s.T(), s.client, &service.APIKey{UserID: userOld.ID, Key: "sk-ul-2", Name: "ul2", Status: service.StatusDisabled})
  604. resetAt := now.Add(10 * time.Minute)
  605. accNormal := mustCreateAccount(s.T(), s.client, &service.Account{Name: "a-normal", Schedulable: true})
  606. mustCreateAccount(s.T(), s.client, &service.Account{Name: "a-error", Status: service.StatusError, Schedulable: true})
  607. mustCreateAccount(s.T(), s.client, &service.Account{Name: "a-rl", RateLimitedAt: &now, RateLimitResetAt: &resetAt, Schedulable: true})
  608. mustCreateAccount(s.T(), s.client, &service.Account{Name: "a-ov", OverloadUntil: &resetAt, Schedulable: true})
  609. d1, d2, d3 := 100, 200, 300
  610. logToday := &service.UsageLog{
  611. UserID: userToday.ID,
  612. APIKeyID: apiKey1.ID,
  613. AccountID: accNormal.ID,
  614. Model: "claude-3",
  615. GroupID: &group.ID,
  616. InputTokens: 10,
  617. OutputTokens: 20,
  618. CacheCreationTokens: 3,
  619. CacheReadTokens: 4,
  620. TotalCost: 1.5,
  621. ActualCost: 1.2,
  622. DurationMs: &d1,
  623. CreatedAt: testMaxTime(todayStart.Add(2*time.Minute), now.Add(-2*time.Minute)),
  624. }
  625. _, err = s.repo.Create(s.ctx, logToday)
  626. s.Require().NoError(err, "Create logToday")
  627. logOld := &service.UsageLog{
  628. UserID: userOld.ID,
  629. APIKeyID: apiKey1.ID,
  630. AccountID: accNormal.ID,
  631. Model: "claude-3",
  632. InputTokens: 5,
  633. OutputTokens: 6,
  634. TotalCost: 0.7,
  635. ActualCost: 0.7,
  636. DurationMs: &d2,
  637. CreatedAt: todayStart.Add(-1 * time.Hour),
  638. }
  639. _, err = s.repo.Create(s.ctx, logOld)
  640. s.Require().NoError(err, "Create logOld")
  641. logPerf := &service.UsageLog{
  642. UserID: userToday.ID,
  643. APIKeyID: apiKey1.ID,
  644. AccountID: accNormal.ID,
  645. Model: "claude-3",
  646. InputTokens: 1,
  647. OutputTokens: 2,
  648. TotalCost: 0.1,
  649. ActualCost: 0.1,
  650. DurationMs: &d3,
  651. CreatedAt: now.Add(-30 * time.Second),
  652. }
  653. _, err = s.repo.Create(s.ctx, logPerf)
  654. s.Require().NoError(err, "Create logPerf")
  655. aggRepo := newDashboardAggregationRepositoryWithSQL(s.tx)
  656. aggStart := todayStart.Add(-2 * time.Hour)
  657. aggEnd := now.Add(2 * time.Minute)
  658. s.Require().NoError(aggRepo.AggregateRange(s.ctx, aggStart, aggEnd), "AggregateRange")
  659. stats, err := s.repo.GetDashboardStats(s.ctx)
  660. s.Require().NoError(err, "GetDashboardStats")
  661. s.Require().Equal(baseStats.TotalUsers+2, stats.TotalUsers, "TotalUsers mismatch")
  662. s.Require().Equal(baseStats.TodayNewUsers+1, stats.TodayNewUsers, "TodayNewUsers mismatch")
  663. s.Require().Equal(baseStats.ActiveUsers+1, stats.ActiveUsers, "ActiveUsers mismatch")
  664. s.Require().Equal(baseStats.TotalAPIKeys+2, stats.TotalAPIKeys, "TotalAPIKeys mismatch")
  665. s.Require().Equal(baseStats.ActiveAPIKeys+1, stats.ActiveAPIKeys, "ActiveAPIKeys mismatch")
  666. s.Require().Equal(baseStats.TotalAccounts+4, stats.TotalAccounts, "TotalAccounts mismatch")
  667. s.Require().Equal(baseStats.ErrorAccounts+1, stats.ErrorAccounts, "ErrorAccounts mismatch")
  668. s.Require().Equal(baseStats.RateLimitAccounts+1, stats.RateLimitAccounts, "RateLimitAccounts mismatch")
  669. s.Require().Equal(baseStats.OverloadAccounts+1, stats.OverloadAccounts, "OverloadAccounts mismatch")
  670. s.Require().Equal(baseStats.TotalRequests+3, stats.TotalRequests, "TotalRequests mismatch")
  671. s.Require().Equal(baseStats.TotalInputTokens+int64(16), stats.TotalInputTokens, "TotalInputTokens mismatch")
  672. s.Require().Equal(baseStats.TotalOutputTokens+int64(28), stats.TotalOutputTokens, "TotalOutputTokens mismatch")
  673. s.Require().Equal(baseStats.TotalCacheCreationTokens+int64(3), stats.TotalCacheCreationTokens, "TotalCacheCreationTokens mismatch")
  674. s.Require().Equal(baseStats.TotalCacheReadTokens+int64(4), stats.TotalCacheReadTokens, "TotalCacheReadTokens mismatch")
  675. s.Require().Equal(baseStats.TotalTokens+int64(51), stats.TotalTokens, "TotalTokens mismatch")
  676. s.Require().Equal(baseStats.TotalCost+2.3, stats.TotalCost, "TotalCost mismatch")
  677. s.Require().Equal(baseStats.TotalActualCost+2.0, stats.TotalActualCost, "TotalActualCost mismatch")
  678. // account_cost falls back to total_cost when account_stats_cost is NULL
  679. s.Require().Equal(baseStats.TotalAccountCost+2.3, stats.TotalAccountCost, "TotalAccountCost mismatch")
  680. s.Require().GreaterOrEqual(stats.TodayRequests, int64(1), "expected TodayRequests >= 1")
  681. s.Require().GreaterOrEqual(stats.TodayCost, 0.0, "expected TodayCost >= 0")
  682. s.Require().GreaterOrEqual(stats.TodayAccountCost, 0.0, "expected TodayAccountCost >= 0")
  683. wantRpm, wantTpm, err := s.repo.getPerformanceStats(s.ctx, 0)
  684. s.Require().NoError(err, "getPerformanceStats")
  685. s.Require().Equal(wantRpm, stats.Rpm, "Rpm mismatch")
  686. s.Require().Equal(wantTpm, stats.Tpm, "Tpm mismatch")
  687. }
  688. func (s *UsageLogRepoSuite) TestDashboardStatsWithRange_Fallback() {
  689. now := time.Now().UTC()
  690. todayStart := truncateToDayUTC(now)
  691. rangeStart := todayStart.Add(-24 * time.Hour)
  692. rangeEnd := now.Add(1 * time.Second)
  693. user1 := mustCreateUser(s.T(), s.client, &service.User{Email: "range-u1@test.com"})
  694. user2 := mustCreateUser(s.T(), s.client, &service.User{Email: "range-u2@test.com"})
  695. apiKey1 := mustCreateApiKey(s.T(), s.client, &service.APIKey{UserID: user1.ID, Key: "sk-range-1", Name: "k1"})
  696. apiKey2 := mustCreateApiKey(s.T(), s.client, &service.APIKey{UserID: user2.ID, Key: "sk-range-2", Name: "k2"})
  697. account := mustCreateAccount(s.T(), s.client, &service.Account{Name: "acc-range"})
  698. d1, d2, d3 := 100, 200, 300
  699. logOutside := &service.UsageLog{
  700. UserID: user1.ID,
  701. APIKeyID: apiKey1.ID,
  702. AccountID: account.ID,
  703. Model: "claude-3",
  704. InputTokens: 7,
  705. OutputTokens: 8,
  706. TotalCost: 0.8,
  707. ActualCost: 0.7,
  708. DurationMs: &d3,
  709. CreatedAt: rangeStart.Add(-1 * time.Hour),
  710. }
  711. _, err := s.repo.Create(s.ctx, logOutside)
  712. s.Require().NoError(err)
  713. logRange := &service.UsageLog{
  714. UserID: user1.ID,
  715. APIKeyID: apiKey1.ID,
  716. AccountID: account.ID,
  717. Model: "claude-3",
  718. InputTokens: 10,
  719. OutputTokens: 20,
  720. CacheCreationTokens: 1,
  721. CacheReadTokens: 2,
  722. TotalCost: 1.0,
  723. ActualCost: 0.9,
  724. DurationMs: &d1,
  725. CreatedAt: rangeStart.Add(2 * time.Hour),
  726. }
  727. _, err = s.repo.Create(s.ctx, logRange)
  728. s.Require().NoError(err)
  729. logToday := &service.UsageLog{
  730. UserID: user2.ID,
  731. APIKeyID: apiKey2.ID,
  732. AccountID: account.ID,
  733. Model: "claude-3",
  734. InputTokens: 5,
  735. OutputTokens: 6,
  736. CacheReadTokens: 1,
  737. TotalCost: 0.5,
  738. ActualCost: 0.5,
  739. DurationMs: &d2,
  740. CreatedAt: now,
  741. }
  742. _, err = s.repo.Create(s.ctx, logToday)
  743. s.Require().NoError(err)
  744. stats, err := s.repo.GetDashboardStatsWithRange(s.ctx, rangeStart, rangeEnd)
  745. s.Require().NoError(err)
  746. s.Require().Equal(int64(2), stats.TotalRequests)
  747. s.Require().Equal(int64(15), stats.TotalInputTokens)
  748. s.Require().Equal(int64(26), stats.TotalOutputTokens)
  749. s.Require().Equal(int64(1), stats.TotalCacheCreationTokens)
  750. s.Require().Equal(int64(3), stats.TotalCacheReadTokens)
  751. s.Require().Equal(int64(45), stats.TotalTokens)
  752. s.Require().Equal(1.5, stats.TotalCost)
  753. s.Require().Equal(1.4, stats.TotalActualCost)
  754. // account_cost = COALESCE(account_stats_cost, total_cost) * COALESCE(account_rate_multiplier, 1) = total_cost
  755. s.Require().Equal(1.5, stats.TotalAccountCost)
  756. s.Require().InEpsilon(150.0, stats.AverageDurationMs, 0.0001)
  757. }
  758. // --- GetUserDashboardStats ---
  759. func (s *UsageLogRepoSuite) TestGetUserDashboardStats() {
  760. user := mustCreateUser(s.T(), s.client, &service.User{Email: "userdash@test.com"})
  761. apiKey := mustCreateApiKey(s.T(), s.client, &service.APIKey{UserID: user.ID, Key: "sk-userdash", Name: "k"})
  762. account := mustCreateAccount(s.T(), s.client, &service.Account{Name: "acc-userdash"})
  763. s.createUsageLog(user, apiKey, account, 10, 20, 0.5, time.Now())
  764. stats, err := s.repo.GetUserDashboardStats(s.ctx, user.ID)
  765. s.Require().NoError(err, "GetUserDashboardStats")
  766. s.Require().Equal(int64(1), stats.TotalAPIKeys)
  767. s.Require().Equal(int64(1), stats.TotalRequests)
  768. }
  769. // --- GetAccountTodayStats ---
  770. func (s *UsageLogRepoSuite) TestGetAccountTodayStats() {
  771. user := mustCreateUser(s.T(), s.client, &service.User{Email: "acctoday@test.com"})
  772. apiKey := mustCreateApiKey(s.T(), s.client, &service.APIKey{UserID: user.ID, Key: "sk-acctoday", Name: "k"})
  773. account := mustCreateAccount(s.T(), s.client, &service.Account{Name: "acc-today"})
  774. createdAt := timezone.Today().Add(1 * time.Hour)
  775. m1 := 1.5
  776. m2 := 0.0
  777. _, err := s.repo.Create(s.ctx, &service.UsageLog{
  778. UserID: user.ID,
  779. APIKeyID: apiKey.ID,
  780. AccountID: account.ID,
  781. RequestID: uuid.New().String(),
  782. Model: "claude-3",
  783. InputTokens: 10,
  784. OutputTokens: 20,
  785. TotalCost: 1.0,
  786. ActualCost: 2.0,
  787. AccountRateMultiplier: &m1,
  788. CreatedAt: createdAt,
  789. })
  790. s.Require().NoError(err)
  791. _, err = s.repo.Create(s.ctx, &service.UsageLog{
  792. UserID: user.ID,
  793. APIKeyID: apiKey.ID,
  794. AccountID: account.ID,
  795. RequestID: uuid.New().String(),
  796. Model: "claude-3",
  797. InputTokens: 5,
  798. OutputTokens: 5,
  799. TotalCost: 0.5,
  800. ActualCost: 1.0,
  801. AccountRateMultiplier: &m2,
  802. CreatedAt: createdAt,
  803. })
  804. s.Require().NoError(err)
  805. stats, err := s.repo.GetAccountTodayStats(s.ctx, account.ID)
  806. s.Require().NoError(err, "GetAccountTodayStats")
  807. s.Require().Equal(int64(2), stats.Requests)
  808. s.Require().Equal(int64(40), stats.Tokens)
  809. // account cost = SUM(total_cost * account_rate_multiplier)
  810. s.Require().InEpsilon(1.5, stats.Cost, 0.0001)
  811. // standard cost = SUM(total_cost)
  812. s.Require().InEpsilon(1.5, stats.StandardCost, 0.0001)
  813. // user cost = SUM(actual_cost)
  814. s.Require().InEpsilon(3.0, stats.UserCost, 0.0001)
  815. }
  816. func (s *UsageLogRepoSuite) TestDashboardAggregationConsistency() {
  817. now := time.Now().UTC().Truncate(time.Second)
  818. // 使用固定的时间偏移确保 hour1 和 hour2 在同一天且都在过去
  819. // 选择当天 02:00 和 03:00 作为测试时间点(基于 now 的日期)
  820. dayStart := truncateToDayUTC(now)
  821. hour1 := dayStart.Add(2 * time.Hour) // 当天 02:00
  822. hour2 := dayStart.Add(3 * time.Hour) // 当天 03:00
  823. // 如果当前时间早于 hour2,则使用昨天的时间
  824. if now.Before(hour2.Add(time.Hour)) {
  825. dayStart = dayStart.Add(-24 * time.Hour)
  826. hour1 = dayStart.Add(2 * time.Hour)
  827. hour2 = dayStart.Add(3 * time.Hour)
  828. }
  829. user1 := mustCreateUser(s.T(), s.client, &service.User{Email: "agg-u1@test.com"})
  830. user2 := mustCreateUser(s.T(), s.client, &service.User{Email: "agg-u2@test.com"})
  831. apiKey1 := mustCreateApiKey(s.T(), s.client, &service.APIKey{UserID: user1.ID, Key: "sk-agg-1", Name: "k1"})
  832. apiKey2 := mustCreateApiKey(s.T(), s.client, &service.APIKey{UserID: user2.ID, Key: "sk-agg-2", Name: "k2"})
  833. account := mustCreateAccount(s.T(), s.client, &service.Account{Name: "acc-agg"})
  834. d1, d2, d3 := 100, 200, 150
  835. log1 := &service.UsageLog{
  836. UserID: user1.ID,
  837. APIKeyID: apiKey1.ID,
  838. AccountID: account.ID,
  839. Model: "claude-3",
  840. InputTokens: 10,
  841. OutputTokens: 20,
  842. CacheCreationTokens: 2,
  843. CacheReadTokens: 1,
  844. TotalCost: 1.0,
  845. ActualCost: 0.9,
  846. DurationMs: &d1,
  847. CreatedAt: hour1.Add(5 * time.Minute),
  848. }
  849. _, err := s.repo.Create(s.ctx, log1)
  850. s.Require().NoError(err)
  851. log2 := &service.UsageLog{
  852. UserID: user1.ID,
  853. APIKeyID: apiKey1.ID,
  854. AccountID: account.ID,
  855. Model: "claude-3",
  856. InputTokens: 5,
  857. OutputTokens: 5,
  858. TotalCost: 0.5,
  859. ActualCost: 0.5,
  860. DurationMs: &d2,
  861. CreatedAt: hour1.Add(20 * time.Minute),
  862. }
  863. _, err = s.repo.Create(s.ctx, log2)
  864. s.Require().NoError(err)
  865. log3 := &service.UsageLog{
  866. UserID: user2.ID,
  867. APIKeyID: apiKey2.ID,
  868. AccountID: account.ID,
  869. Model: "claude-3",
  870. InputTokens: 7,
  871. OutputTokens: 8,
  872. TotalCost: 0.7,
  873. ActualCost: 0.7,
  874. DurationMs: &d3,
  875. CreatedAt: hour2.Add(10 * time.Minute),
  876. }
  877. _, err = s.repo.Create(s.ctx, log3)
  878. s.Require().NoError(err)
  879. aggRepo := newDashboardAggregationRepositoryWithSQL(s.tx)
  880. aggStart := hour1.Add(-5 * time.Minute)
  881. aggEnd := hour2.Add(time.Hour) // 确保覆盖 hour2 的所有数据
  882. s.Require().NoError(aggRepo.AggregateRange(s.ctx, aggStart, aggEnd))
  883. type hourlyRow struct {
  884. totalRequests int64
  885. inputTokens int64
  886. outputTokens int64
  887. cacheCreationTokens int64
  888. cacheReadTokens int64
  889. totalCost float64
  890. actualCost float64
  891. totalDurationMs int64
  892. activeUsers int64
  893. }
  894. fetchHourly := func(bucketStart time.Time) hourlyRow {
  895. var row hourlyRow
  896. err := scanSingleRow(s.ctx, s.tx, `
  897. SELECT total_requests, input_tokens, output_tokens, cache_creation_tokens, cache_read_tokens,
  898. total_cost, actual_cost, total_duration_ms, active_users
  899. FROM usage_dashboard_hourly
  900. WHERE bucket_start = $1
  901. `, []any{bucketStart}, &row.totalRequests, &row.inputTokens, &row.outputTokens,
  902. &row.cacheCreationTokens, &row.cacheReadTokens, &row.totalCost, &row.actualCost,
  903. &row.totalDurationMs, &row.activeUsers,
  904. )
  905. s.Require().NoError(err)
  906. return row
  907. }
  908. hour1Row := fetchHourly(hour1)
  909. s.Require().Equal(int64(2), hour1Row.totalRequests)
  910. s.Require().Equal(int64(15), hour1Row.inputTokens)
  911. s.Require().Equal(int64(25), hour1Row.outputTokens)
  912. s.Require().Equal(int64(2), hour1Row.cacheCreationTokens)
  913. s.Require().Equal(int64(1), hour1Row.cacheReadTokens)
  914. s.Require().Equal(1.5, hour1Row.totalCost)
  915. s.Require().Equal(1.4, hour1Row.actualCost)
  916. s.Require().Equal(int64(300), hour1Row.totalDurationMs)
  917. s.Require().Equal(int64(1), hour1Row.activeUsers)
  918. hour2Row := fetchHourly(hour2)
  919. s.Require().Equal(int64(1), hour2Row.totalRequests)
  920. s.Require().Equal(int64(7), hour2Row.inputTokens)
  921. s.Require().Equal(int64(8), hour2Row.outputTokens)
  922. s.Require().Equal(int64(0), hour2Row.cacheCreationTokens)
  923. s.Require().Equal(int64(0), hour2Row.cacheReadTokens)
  924. s.Require().Equal(0.7, hour2Row.totalCost)
  925. s.Require().Equal(0.7, hour2Row.actualCost)
  926. s.Require().Equal(int64(150), hour2Row.totalDurationMs)
  927. s.Require().Equal(int64(1), hour2Row.activeUsers)
  928. var daily struct {
  929. totalRequests int64
  930. inputTokens int64
  931. outputTokens int64
  932. cacheCreationTokens int64
  933. cacheReadTokens int64
  934. totalCost float64
  935. actualCost float64
  936. totalDurationMs int64
  937. activeUsers int64
  938. }
  939. err = scanSingleRow(s.ctx, s.tx, `
  940. SELECT total_requests, input_tokens, output_tokens, cache_creation_tokens, cache_read_tokens,
  941. total_cost, actual_cost, total_duration_ms, active_users
  942. FROM usage_dashboard_daily
  943. WHERE bucket_date = $1::date
  944. `, []any{dayStart}, &daily.totalRequests, &daily.inputTokens, &daily.outputTokens,
  945. &daily.cacheCreationTokens, &daily.cacheReadTokens, &daily.totalCost, &daily.actualCost,
  946. &daily.totalDurationMs, &daily.activeUsers,
  947. )
  948. s.Require().NoError(err)
  949. s.Require().Equal(int64(3), daily.totalRequests)
  950. s.Require().Equal(int64(22), daily.inputTokens)
  951. s.Require().Equal(int64(33), daily.outputTokens)
  952. s.Require().Equal(int64(2), daily.cacheCreationTokens)
  953. s.Require().Equal(int64(1), daily.cacheReadTokens)
  954. s.Require().Equal(2.2, daily.totalCost)
  955. s.Require().Equal(2.1, daily.actualCost)
  956. s.Require().Equal(int64(450), daily.totalDurationMs)
  957. s.Require().Equal(int64(2), daily.activeUsers)
  958. }
  959. // --- GetBatchUserUsageStats ---
  960. func (s *UsageLogRepoSuite) TestGetBatchUserUsageStats() {
  961. user1 := mustCreateUser(s.T(), s.client, &service.User{Email: "batch1@test.com"})
  962. user2 := mustCreateUser(s.T(), s.client, &service.User{Email: "batch2@test.com"})
  963. apiKey1 := mustCreateApiKey(s.T(), s.client, &service.APIKey{UserID: user1.ID, Key: "sk-batch1", Name: "k"})
  964. apiKey2 := mustCreateApiKey(s.T(), s.client, &service.APIKey{UserID: user2.ID, Key: "sk-batch2", Name: "k"})
  965. account := mustCreateAccount(s.T(), s.client, &service.Account{Name: "acc-batch"})
  966. s.createUsageLog(user1, apiKey1, account, 10, 20, 0.5, time.Now())
  967. s.createUsageLog(user2, apiKey2, account, 15, 25, 0.6, time.Now())
  968. stats, err := s.repo.GetBatchUserUsageStats(s.ctx, []int64{user1.ID, user2.ID}, time.Time{}, time.Time{})
  969. s.Require().NoError(err, "GetBatchUserUsageStats")
  970. s.Require().Len(stats, 2)
  971. s.Require().NotNil(stats[user1.ID])
  972. s.Require().NotNil(stats[user2.ID])
  973. }
  974. func (s *UsageLogRepoSuite) TestGetBatchUserUsageStats_Empty() {
  975. stats, err := s.repo.GetBatchUserUsageStats(s.ctx, []int64{}, time.Time{}, time.Time{})
  976. s.Require().NoError(err)
  977. s.Require().Empty(stats)
  978. }
  979. // --- GetBatchAPIKeyUsageStats ---
  980. func (s *UsageLogRepoSuite) TestGetBatchApiKeyUsageStats() {
  981. user := mustCreateUser(s.T(), s.client, &service.User{Email: "batchkey@test.com"})
  982. apiKey1 := mustCreateApiKey(s.T(), s.client, &service.APIKey{UserID: user.ID, Key: "sk-batchkey1", Name: "k1"})
  983. apiKey2 := mustCreateApiKey(s.T(), s.client, &service.APIKey{UserID: user.ID, Key: "sk-batchkey2", Name: "k2"})
  984. account := mustCreateAccount(s.T(), s.client, &service.Account{Name: "acc-batchkey"})
  985. s.createUsageLog(user, apiKey1, account, 10, 20, 0.5, time.Now())
  986. s.createUsageLog(user, apiKey2, account, 15, 25, 0.6, time.Now())
  987. stats, err := s.repo.GetBatchAPIKeyUsageStats(s.ctx, []int64{apiKey1.ID, apiKey2.ID}, time.Time{}, time.Time{})
  988. s.Require().NoError(err, "GetBatchAPIKeyUsageStats")
  989. s.Require().Len(stats, 2)
  990. }
  991. func (s *UsageLogRepoSuite) TestGetBatchApiKeyUsageStats_Empty() {
  992. stats, err := s.repo.GetBatchAPIKeyUsageStats(s.ctx, []int64{}, time.Time{}, time.Time{})
  993. s.Require().NoError(err)
  994. s.Require().Empty(stats)
  995. }
  996. // --- GetGlobalStats ---
  997. func (s *UsageLogRepoSuite) TestGetGlobalStats() {
  998. user := mustCreateUser(s.T(), s.client, &service.User{Email: "global@test.com"})
  999. apiKey := mustCreateApiKey(s.T(), s.client, &service.APIKey{UserID: user.ID, Key: "sk-global", Name: "k"})
  1000. account := mustCreateAccount(s.T(), s.client, &service.Account{Name: "acc-global"})
  1001. base := time.Date(2025, 1, 15, 12, 0, 0, 0, time.UTC)
  1002. s.createUsageLog(user, apiKey, account, 10, 20, 0.5, base)
  1003. s.createUsageLog(user, apiKey, account, 15, 25, 0.6, base.Add(1*time.Hour))
  1004. stats, err := s.repo.GetGlobalStats(s.ctx, base.Add(-1*time.Hour), base.Add(2*time.Hour))
  1005. s.Require().NoError(err, "GetGlobalStats")
  1006. s.Require().Equal(int64(2), stats.TotalRequests)
  1007. s.Require().Equal(int64(25), stats.TotalInputTokens)
  1008. s.Require().Equal(int64(45), stats.TotalOutputTokens)
  1009. }
  1010. func testMaxTime(a, b time.Time) time.Time {
  1011. if a.After(b) {
  1012. return a
  1013. }
  1014. return b
  1015. }
  1016. // --- ListByUserAndTimeRange ---
  1017. func (s *UsageLogRepoSuite) TestListByUserAndTimeRange() {
  1018. user := mustCreateUser(s.T(), s.client, &service.User{Email: "timerange@test.com"})
  1019. apiKey := mustCreateApiKey(s.T(), s.client, &service.APIKey{UserID: user.ID, Key: "sk-timerange", Name: "k"})
  1020. account := mustCreateAccount(s.T(), s.client, &service.Account{Name: "acc-timerange"})
  1021. base := time.Date(2025, 1, 15, 12, 0, 0, 0, time.UTC)
  1022. s.createUsageLog(user, apiKey, account, 10, 20, 0.5, base)
  1023. s.createUsageLog(user, apiKey, account, 15, 25, 0.6, base.Add(1*time.Hour))
  1024. s.createUsageLog(user, apiKey, account, 20, 30, 0.7, base.Add(-24*time.Hour)) // outside range
  1025. startTime := base.Add(-1 * time.Hour)
  1026. endTime := base.Add(2 * time.Hour)
  1027. logs, _, err := s.repo.ListByUserAndTimeRange(s.ctx, user.ID, startTime, endTime)
  1028. s.Require().NoError(err, "ListByUserAndTimeRange")
  1029. s.Require().Len(logs, 2)
  1030. }
  1031. // --- ListByAPIKeyAndTimeRange ---
  1032. func (s *UsageLogRepoSuite) TestListByAPIKeyAndTimeRange() {
  1033. user := mustCreateUser(s.T(), s.client, &service.User{Email: "keytimerange@test.com"})
  1034. apiKey := mustCreateApiKey(s.T(), s.client, &service.APIKey{UserID: user.ID, Key: "sk-keytimerange", Name: "k"})
  1035. account := mustCreateAccount(s.T(), s.client, &service.Account{Name: "acc-keytimerange"})
  1036. base := time.Date(2025, 1, 15, 12, 0, 0, 0, time.UTC)
  1037. s.createUsageLog(user, apiKey, account, 10, 20, 0.5, base)
  1038. s.createUsageLog(user, apiKey, account, 15, 25, 0.6, base.Add(30*time.Minute))
  1039. s.createUsageLog(user, apiKey, account, 20, 30, 0.7, base.Add(-24*time.Hour)) // outside range
  1040. startTime := base.Add(-1 * time.Hour)
  1041. endTime := base.Add(2 * time.Hour)
  1042. logs, _, err := s.repo.ListByAPIKeyAndTimeRange(s.ctx, apiKey.ID, startTime, endTime)
  1043. s.Require().NoError(err, "ListByAPIKeyAndTimeRange")
  1044. s.Require().Len(logs, 2)
  1045. }
  1046. // --- ListByAccountAndTimeRange ---
  1047. func (s *UsageLogRepoSuite) TestListByAccountAndTimeRange() {
  1048. user := mustCreateUser(s.T(), s.client, &service.User{Email: "acctimerange@test.com"})
  1049. apiKey := mustCreateApiKey(s.T(), s.client, &service.APIKey{UserID: user.ID, Key: "sk-acctimerange", Name: "k"})
  1050. account := mustCreateAccount(s.T(), s.client, &service.Account{Name: "acc-acctimerange"})
  1051. base := time.Date(2025, 1, 15, 12, 0, 0, 0, time.UTC)
  1052. s.createUsageLog(user, apiKey, account, 10, 20, 0.5, base)
  1053. s.createUsageLog(user, apiKey, account, 15, 25, 0.6, base.Add(45*time.Minute))
  1054. s.createUsageLog(user, apiKey, account, 20, 30, 0.7, base.Add(-24*time.Hour)) // outside range
  1055. startTime := base.Add(-1 * time.Hour)
  1056. endTime := base.Add(2 * time.Hour)
  1057. logs, _, err := s.repo.ListByAccountAndTimeRange(s.ctx, account.ID, startTime, endTime)
  1058. s.Require().NoError(err, "ListByAccountAndTimeRange")
  1059. s.Require().Len(logs, 2)
  1060. }
  1061. // --- ListByModelAndTimeRange ---
  1062. func (s *UsageLogRepoSuite) TestListByModelAndTimeRange() {
  1063. user := mustCreateUser(s.T(), s.client, &service.User{Email: "modeltimerange@test.com"})
  1064. apiKey := mustCreateApiKey(s.T(), s.client, &service.APIKey{UserID: user.ID, Key: "sk-modeltimerange", Name: "k"})
  1065. account := mustCreateAccount(s.T(), s.client, &service.Account{Name: "acc-modeltimerange"})
  1066. base := time.Date(2025, 1, 15, 12, 0, 0, 0, time.UTC)
  1067. // Create logs with different models
  1068. log1 := &service.UsageLog{
  1069. UserID: user.ID,
  1070. APIKeyID: apiKey.ID,
  1071. AccountID: account.ID,
  1072. Model: "claude-3-opus",
  1073. InputTokens: 10,
  1074. OutputTokens: 20,
  1075. TotalCost: 0.5,
  1076. ActualCost: 0.5,
  1077. CreatedAt: base,
  1078. }
  1079. _, err := s.repo.Create(s.ctx, log1)
  1080. s.Require().NoError(err)
  1081. log2 := &service.UsageLog{
  1082. UserID: user.ID,
  1083. APIKeyID: apiKey.ID,
  1084. AccountID: account.ID,
  1085. Model: "claude-3-opus",
  1086. InputTokens: 15,
  1087. OutputTokens: 25,
  1088. TotalCost: 0.6,
  1089. ActualCost: 0.6,
  1090. CreatedAt: base.Add(30 * time.Minute),
  1091. }
  1092. _, err = s.repo.Create(s.ctx, log2)
  1093. s.Require().NoError(err)
  1094. log3 := &service.UsageLog{
  1095. UserID: user.ID,
  1096. APIKeyID: apiKey.ID,
  1097. AccountID: account.ID,
  1098. Model: "claude-3-sonnet",
  1099. InputTokens: 20,
  1100. OutputTokens: 30,
  1101. TotalCost: 0.7,
  1102. ActualCost: 0.7,
  1103. CreatedAt: base.Add(1 * time.Hour),
  1104. }
  1105. _, err = s.repo.Create(s.ctx, log3)
  1106. s.Require().NoError(err)
  1107. startTime := base.Add(-1 * time.Hour)
  1108. endTime := base.Add(2 * time.Hour)
  1109. logs, _, err := s.repo.ListByModelAndTimeRange(s.ctx, "claude-3-opus", startTime, endTime)
  1110. s.Require().NoError(err, "ListByModelAndTimeRange")
  1111. s.Require().Len(logs, 2)
  1112. }
  1113. // --- GetAccountWindowStats ---
  1114. func (s *UsageLogRepoSuite) TestGetAccountWindowStats() {
  1115. user := mustCreateUser(s.T(), s.client, &service.User{Email: "windowstats@test.com"})
  1116. apiKey := mustCreateApiKey(s.T(), s.client, &service.APIKey{UserID: user.ID, Key: "sk-windowstats", Name: "k"})
  1117. account := mustCreateAccount(s.T(), s.client, &service.Account{Name: "acc-windowstats"})
  1118. now := time.Now()
  1119. windowStart := now.Add(-10 * time.Minute)
  1120. s.createUsageLog(user, apiKey, account, 10, 20, 0.5, now.Add(-5*time.Minute))
  1121. s.createUsageLog(user, apiKey, account, 15, 25, 0.6, now.Add(-3*time.Minute))
  1122. s.createUsageLog(user, apiKey, account, 20, 30, 0.7, now.Add(-30*time.Minute)) // outside window
  1123. stats, err := s.repo.GetAccountWindowStats(s.ctx, account.ID, windowStart)
  1124. s.Require().NoError(err, "GetAccountWindowStats")
  1125. s.Require().Equal(int64(2), stats.Requests)
  1126. s.Require().Equal(int64(70), stats.Tokens) // (10+20) + (15+25)
  1127. }
  1128. // --- GetUserUsageTrendByUserID ---
  1129. func (s *UsageLogRepoSuite) TestGetUserUsageTrendByUserID() {
  1130. user := mustCreateUser(s.T(), s.client, &service.User{Email: "usertrend@test.com"})
  1131. apiKey := mustCreateApiKey(s.T(), s.client, &service.APIKey{UserID: user.ID, Key: "sk-usertrend", Name: "k"})
  1132. account := mustCreateAccount(s.T(), s.client, &service.Account{Name: "acc-usertrend"})
  1133. base := time.Date(2025, 1, 15, 12, 0, 0, 0, time.UTC)
  1134. s.createUsageLog(user, apiKey, account, 10, 20, 0.5, base)
  1135. s.createUsageLog(user, apiKey, account, 15, 25, 0.6, base.Add(1*time.Hour))
  1136. s.createUsageLog(user, apiKey, account, 20, 30, 0.7, base.Add(24*time.Hour)) // next day
  1137. startTime := base.Add(-1 * time.Hour)
  1138. endTime := base.Add(48 * time.Hour)
  1139. trend, err := s.repo.GetUserUsageTrendByUserID(s.ctx, user.ID, startTime, endTime, "day")
  1140. s.Require().NoError(err, "GetUserUsageTrendByUserID")
  1141. s.Require().Len(trend, 2) // 2 different days
  1142. }
  1143. func (s *UsageLogRepoSuite) TestGetUserUsageTrendByUserID_HourlyGranularity() {
  1144. user := mustCreateUser(s.T(), s.client, &service.User{Email: "usertrendhourly@test.com"})
  1145. apiKey := mustCreateApiKey(s.T(), s.client, &service.APIKey{UserID: user.ID, Key: "sk-usertrendhourly", Name: "k"})
  1146. account := mustCreateAccount(s.T(), s.client, &service.Account{Name: "acc-usertrendhourly"})
  1147. base := time.Date(2025, 1, 15, 12, 0, 0, 0, time.UTC)
  1148. s.createUsageLog(user, apiKey, account, 10, 20, 0.5, base)
  1149. s.createUsageLog(user, apiKey, account, 15, 25, 0.6, base.Add(1*time.Hour))
  1150. s.createUsageLog(user, apiKey, account, 20, 30, 0.7, base.Add(2*time.Hour))
  1151. startTime := base.Add(-1 * time.Hour)
  1152. endTime := base.Add(3 * time.Hour)
  1153. trend, err := s.repo.GetUserUsageTrendByUserID(s.ctx, user.ID, startTime, endTime, "hour")
  1154. s.Require().NoError(err, "GetUserUsageTrendByUserID hourly")
  1155. s.Require().Len(trend, 3) // 3 different hours
  1156. }
  1157. // --- GetUserModelStats ---
  1158. func (s *UsageLogRepoSuite) TestGetUserModelStats() {
  1159. user := mustCreateUser(s.T(), s.client, &service.User{Email: "modelstats@test.com"})
  1160. apiKey := mustCreateApiKey(s.T(), s.client, &service.APIKey{UserID: user.ID, Key: "sk-modelstats", Name: "k"})
  1161. account := mustCreateAccount(s.T(), s.client, &service.Account{Name: "acc-modelstats"})
  1162. base := time.Date(2025, 1, 15, 12, 0, 0, 0, time.UTC)
  1163. // Create logs with different models
  1164. log1 := &service.UsageLog{
  1165. UserID: user.ID,
  1166. APIKeyID: apiKey.ID,
  1167. AccountID: account.ID,
  1168. Model: "claude-3-opus",
  1169. InputTokens: 100,
  1170. OutputTokens: 200,
  1171. TotalCost: 0.5,
  1172. ActualCost: 0.5,
  1173. CreatedAt: base,
  1174. }
  1175. _, err := s.repo.Create(s.ctx, log1)
  1176. s.Require().NoError(err)
  1177. log2 := &service.UsageLog{
  1178. UserID: user.ID,
  1179. APIKeyID: apiKey.ID,
  1180. AccountID: account.ID,
  1181. Model: "claude-3-sonnet",
  1182. InputTokens: 50,
  1183. OutputTokens: 100,
  1184. TotalCost: 0.2,
  1185. ActualCost: 0.2,
  1186. CreatedAt: base.Add(1 * time.Hour),
  1187. }
  1188. _, err = s.repo.Create(s.ctx, log2)
  1189. s.Require().NoError(err)
  1190. startTime := base.Add(-1 * time.Hour)
  1191. endTime := base.Add(2 * time.Hour)
  1192. stats, err := s.repo.GetUserModelStats(s.ctx, user.ID, startTime, endTime)
  1193. s.Require().NoError(err, "GetUserModelStats")
  1194. s.Require().Len(stats, 2)
  1195. // Should be ordered by total_tokens DESC
  1196. s.Require().Equal("claude-3-opus", stats[0].Model)
  1197. s.Require().Equal(int64(300), stats[0].TotalTokens)
  1198. }
  1199. // --- GetUsageTrendWithFilters ---
  1200. func (s *UsageLogRepoSuite) TestGetUsageTrendWithFilters() {
  1201. user := mustCreateUser(s.T(), s.client, &service.User{Email: "trendfilters@test.com"})
  1202. apiKey := mustCreateApiKey(s.T(), s.client, &service.APIKey{UserID: user.ID, Key: "sk-trendfilters", Name: "k"})
  1203. account := mustCreateAccount(s.T(), s.client, &service.Account{Name: "acc-trendfilters"})
  1204. base := time.Date(2025, 1, 15, 12, 0, 0, 0, time.UTC)
  1205. s.createUsageLog(user, apiKey, account, 10, 20, 0.5, base)
  1206. s.createUsageLog(user, apiKey, account, 15, 25, 0.6, base.Add(24*time.Hour))
  1207. startTime := base.Add(-1 * time.Hour)
  1208. endTime := base.Add(48 * time.Hour)
  1209. // Test with user filter
  1210. trend, err := s.repo.GetUsageTrendWithFilters(s.ctx, startTime, endTime, "day", user.ID, 0, 0, 0, "", nil, nil, nil)
  1211. s.Require().NoError(err, "GetUsageTrendWithFilters user filter")
  1212. s.Require().Len(trend, 2)
  1213. // Test with apiKey filter
  1214. trend, err = s.repo.GetUsageTrendWithFilters(s.ctx, startTime, endTime, "day", 0, apiKey.ID, 0, 0, "", nil, nil, nil)
  1215. s.Require().NoError(err, "GetUsageTrendWithFilters apiKey filter")
  1216. s.Require().Len(trend, 2)
  1217. // Test with both filters
  1218. trend, err = s.repo.GetUsageTrendWithFilters(s.ctx, startTime, endTime, "day", user.ID, apiKey.ID, 0, 0, "", nil, nil, nil)
  1219. s.Require().NoError(err, "GetUsageTrendWithFilters both filters")
  1220. s.Require().Len(trend, 2)
  1221. }
  1222. func (s *UsageLogRepoSuite) TestGetUsageTrendWithFilters_HourlyGranularity() {
  1223. user := mustCreateUser(s.T(), s.client, &service.User{Email: "trendfilters-h@test.com"})
  1224. apiKey := mustCreateApiKey(s.T(), s.client, &service.APIKey{UserID: user.ID, Key: "sk-trendfilters-h", Name: "k"})
  1225. account := mustCreateAccount(s.T(), s.client, &service.Account{Name: "acc-trendfilters-h"})
  1226. base := time.Date(2025, 1, 15, 12, 0, 0, 0, time.UTC)
  1227. s.createUsageLog(user, apiKey, account, 10, 20, 0.5, base)
  1228. s.createUsageLog(user, apiKey, account, 15, 25, 0.6, base.Add(1*time.Hour))
  1229. startTime := base.Add(-1 * time.Hour)
  1230. endTime := base.Add(3 * time.Hour)
  1231. trend, err := s.repo.GetUsageTrendWithFilters(s.ctx, startTime, endTime, "hour", user.ID, 0, 0, 0, "", nil, nil, nil)
  1232. s.Require().NoError(err, "GetUsageTrendWithFilters hourly")
  1233. s.Require().Len(trend, 2)
  1234. }
  1235. // --- GetModelStatsWithFilters ---
  1236. func (s *UsageLogRepoSuite) TestGetModelStatsWithFilters() {
  1237. user := mustCreateUser(s.T(), s.client, &service.User{Email: "modelfilters@test.com"})
  1238. apiKey := mustCreateApiKey(s.T(), s.client, &service.APIKey{UserID: user.ID, Key: "sk-modelfilters", Name: "k"})
  1239. account := mustCreateAccount(s.T(), s.client, &service.Account{Name: "acc-modelfilters"})
  1240. base := time.Date(2025, 1, 15, 12, 0, 0, 0, time.UTC)
  1241. log1 := &service.UsageLog{
  1242. UserID: user.ID,
  1243. APIKeyID: apiKey.ID,
  1244. AccountID: account.ID,
  1245. Model: "claude-3-opus",
  1246. InputTokens: 100,
  1247. OutputTokens: 200,
  1248. TotalCost: 0.5,
  1249. ActualCost: 0.5,
  1250. CreatedAt: base,
  1251. }
  1252. _, err := s.repo.Create(s.ctx, log1)
  1253. s.Require().NoError(err)
  1254. log2 := &service.UsageLog{
  1255. UserID: user.ID,
  1256. APIKeyID: apiKey.ID,
  1257. AccountID: account.ID,
  1258. Model: "claude-3-sonnet",
  1259. InputTokens: 50,
  1260. OutputTokens: 100,
  1261. TotalCost: 0.2,
  1262. ActualCost: 0.2,
  1263. CreatedAt: base.Add(1 * time.Hour),
  1264. }
  1265. _, err = s.repo.Create(s.ctx, log2)
  1266. s.Require().NoError(err)
  1267. startTime := base.Add(-1 * time.Hour)
  1268. endTime := base.Add(2 * time.Hour)
  1269. // Test with user filter
  1270. stats, err := s.repo.GetModelStatsWithFilters(s.ctx, startTime, endTime, user.ID, 0, 0, 0, nil, nil, nil)
  1271. s.Require().NoError(err, "GetModelStatsWithFilters user filter")
  1272. s.Require().Len(stats, 2)
  1273. // Test with apiKey filter
  1274. stats, err = s.repo.GetModelStatsWithFilters(s.ctx, startTime, endTime, 0, apiKey.ID, 0, 0, nil, nil, nil)
  1275. s.Require().NoError(err, "GetModelStatsWithFilters apiKey filter")
  1276. s.Require().Len(stats, 2)
  1277. // Test with account filter
  1278. stats, err = s.repo.GetModelStatsWithFilters(s.ctx, startTime, endTime, 0, 0, account.ID, 0, nil, nil, nil)
  1279. s.Require().NoError(err, "GetModelStatsWithFilters account filter")
  1280. s.Require().Len(stats, 2)
  1281. }
  1282. // --- GetAccountUsageStats ---
  1283. func (s *UsageLogRepoSuite) TestGetAccountUsageStats() {
  1284. user := mustCreateUser(s.T(), s.client, &service.User{Email: "accstats@test.com"})
  1285. apiKey := mustCreateApiKey(s.T(), s.client, &service.APIKey{UserID: user.ID, Key: "sk-accstats", Name: "k"})
  1286. account := mustCreateAccount(s.T(), s.client, &service.Account{Name: "acc-accstats"})
  1287. base := time.Date(2025, 1, 15, 0, 0, 0, 0, time.UTC)
  1288. // Create logs on different days
  1289. log1 := &service.UsageLog{
  1290. UserID: user.ID,
  1291. APIKeyID: apiKey.ID,
  1292. AccountID: account.ID,
  1293. Model: "claude-3-opus",
  1294. InputTokens: 100,
  1295. OutputTokens: 200,
  1296. TotalCost: 0.5,
  1297. ActualCost: 0.4,
  1298. CreatedAt: base.Add(12 * time.Hour),
  1299. }
  1300. _, err := s.repo.Create(s.ctx, log1)
  1301. s.Require().NoError(err)
  1302. log2 := &service.UsageLog{
  1303. UserID: user.ID,
  1304. APIKeyID: apiKey.ID,
  1305. AccountID: account.ID,
  1306. Model: "claude-3-sonnet",
  1307. InputTokens: 50,
  1308. OutputTokens: 100,
  1309. TotalCost: 0.2,
  1310. ActualCost: 0.15,
  1311. CreatedAt: base.Add(36 * time.Hour), // next day
  1312. }
  1313. _, err = s.repo.Create(s.ctx, log2)
  1314. s.Require().NoError(err)
  1315. startTime := base
  1316. endTime := base.Add(72 * time.Hour)
  1317. resp, err := s.repo.GetAccountUsageStats(s.ctx, account.ID, startTime, endTime)
  1318. s.Require().NoError(err, "GetAccountUsageStats")
  1319. s.Require().Len(resp.History, 2, "expected 2 days of history")
  1320. s.Require().Equal(int64(2), resp.Summary.TotalRequests)
  1321. s.Require().Equal(int64(450), resp.Summary.TotalTokens)
  1322. s.Require().Len(resp.Models, 2)
  1323. }
  1324. func (s *UsageLogRepoSuite) TestGetAccountUsageStats_EmptyRange() {
  1325. account := mustCreateAccount(s.T(), s.client, &service.Account{Name: "acc-emptystats"})
  1326. base := time.Date(2025, 1, 15, 0, 0, 0, 0, time.UTC)
  1327. startTime := base
  1328. endTime := base.Add(72 * time.Hour)
  1329. resp, err := s.repo.GetAccountUsageStats(s.ctx, account.ID, startTime, endTime)
  1330. s.Require().NoError(err, "GetAccountUsageStats empty")
  1331. s.Require().Len(resp.History, 0)
  1332. s.Require().Equal(int64(0), resp.Summary.TotalRequests)
  1333. }
  1334. // --- GetUserUsageTrend ---
  1335. func (s *UsageLogRepoSuite) TestGetUserUsageTrend() {
  1336. user1 := mustCreateUser(s.T(), s.client, &service.User{Email: "usertrend1@test.com"})
  1337. user2 := mustCreateUser(s.T(), s.client, &service.User{Email: "usertrend2@test.com"})
  1338. apiKey1 := mustCreateApiKey(s.T(), s.client, &service.APIKey{UserID: user1.ID, Key: "sk-usertrend1", Name: "k1"})
  1339. apiKey2 := mustCreateApiKey(s.T(), s.client, &service.APIKey{UserID: user2.ID, Key: "sk-usertrend2", Name: "k2"})
  1340. account := mustCreateAccount(s.T(), s.client, &service.Account{Name: "acc-usertrends"})
  1341. base := time.Date(2025, 1, 15, 12, 0, 0, 0, time.UTC)
  1342. s.createUsageLog(user1, apiKey1, account, 100, 200, 1.0, base)
  1343. s.createUsageLog(user2, apiKey2, account, 50, 100, 0.5, base)
  1344. s.createUsageLog(user1, apiKey1, account, 100, 200, 1.0, base.Add(24*time.Hour))
  1345. startTime := base.Add(-1 * time.Hour)
  1346. endTime := base.Add(48 * time.Hour)
  1347. trend, err := s.repo.GetUserUsageTrend(s.ctx, startTime, endTime, "day", 10)
  1348. s.Require().NoError(err, "GetUserUsageTrend")
  1349. s.Require().GreaterOrEqual(len(trend), 2)
  1350. }
  1351. // --- GetAPIKeyUsageTrend ---
  1352. func (s *UsageLogRepoSuite) TestGetAPIKeyUsageTrend() {
  1353. user := mustCreateUser(s.T(), s.client, &service.User{Email: "keytrend@test.com"})
  1354. apiKey1 := mustCreateApiKey(s.T(), s.client, &service.APIKey{UserID: user.ID, Key: "sk-keytrend1", Name: "k1"})
  1355. apiKey2 := mustCreateApiKey(s.T(), s.client, &service.APIKey{UserID: user.ID, Key: "sk-keytrend2", Name: "k2"})
  1356. account := mustCreateAccount(s.T(), s.client, &service.Account{Name: "acc-keytrends"})
  1357. base := time.Date(2025, 1, 15, 12, 0, 0, 0, time.UTC)
  1358. s.createUsageLog(user, apiKey1, account, 100, 200, 1.0, base)
  1359. s.createUsageLog(user, apiKey2, account, 50, 100, 0.5, base)
  1360. s.createUsageLog(user, apiKey1, account, 100, 200, 1.0, base.Add(24*time.Hour))
  1361. startTime := base.Add(-1 * time.Hour)
  1362. endTime := base.Add(48 * time.Hour)
  1363. trend, err := s.repo.GetAPIKeyUsageTrend(s.ctx, startTime, endTime, "day", 10)
  1364. s.Require().NoError(err, "GetAPIKeyUsageTrend")
  1365. s.Require().GreaterOrEqual(len(trend), 2)
  1366. }
  1367. func (s *UsageLogRepoSuite) TestGetAPIKeyUsageTrend_HourlyGranularity() {
  1368. user := mustCreateUser(s.T(), s.client, &service.User{Email: "keytrendh@test.com"})
  1369. apiKey := mustCreateApiKey(s.T(), s.client, &service.APIKey{UserID: user.ID, Key: "sk-keytrendh", Name: "k"})
  1370. account := mustCreateAccount(s.T(), s.client, &service.Account{Name: "acc-keytrendh"})
  1371. base := time.Date(2025, 1, 15, 12, 0, 0, 0, time.UTC)
  1372. s.createUsageLog(user, apiKey, account, 100, 200, 1.0, base)
  1373. s.createUsageLog(user, apiKey, account, 50, 100, 0.5, base.Add(1*time.Hour))
  1374. startTime := base.Add(-1 * time.Hour)
  1375. endTime := base.Add(3 * time.Hour)
  1376. trend, err := s.repo.GetAPIKeyUsageTrend(s.ctx, startTime, endTime, "hour", 10)
  1377. s.Require().NoError(err, "GetAPIKeyUsageTrend hourly")
  1378. s.Require().Len(trend, 2)
  1379. }
  1380. // --- ListWithFilters (additional filter tests) ---
  1381. func (s *UsageLogRepoSuite) TestListWithFilters_ApiKeyFilter() {
  1382. user := mustCreateUser(s.T(), s.client, &service.User{Email: "filterskey@test.com"})
  1383. apiKey := mustCreateApiKey(s.T(), s.client, &service.APIKey{UserID: user.ID, Key: "sk-filterskey", Name: "k"})
  1384. account := mustCreateAccount(s.T(), s.client, &service.Account{Name: "acc-filterskey"})
  1385. s.createUsageLog(user, apiKey, account, 10, 20, 0.5, time.Now())
  1386. filters := usagestats.UsageLogFilters{APIKeyID: apiKey.ID}
  1387. logs, page, err := s.repo.ListWithFilters(s.ctx, pagination.PaginationParams{Page: 1, PageSize: 10}, filters)
  1388. s.Require().NoError(err, "ListWithFilters apiKey")
  1389. s.Require().Len(logs, 1)
  1390. s.Require().Equal(int64(1), page.Total)
  1391. }
  1392. func (s *UsageLogRepoSuite) TestListWithFilters_TimeRange() {
  1393. user := mustCreateUser(s.T(), s.client, &service.User{Email: "filterstime@test.com"})
  1394. apiKey := mustCreateApiKey(s.T(), s.client, &service.APIKey{UserID: user.ID, Key: "sk-filterstime", Name: "k"})
  1395. account := mustCreateAccount(s.T(), s.client, &service.Account{Name: "acc-filterstime"})
  1396. base := time.Date(2025, 1, 15, 12, 0, 0, 0, time.UTC)
  1397. s.createUsageLog(user, apiKey, account, 10, 20, 0.5, base)
  1398. s.createUsageLog(user, apiKey, account, 15, 25, 0.6, base.Add(1*time.Hour))
  1399. s.createUsageLog(user, apiKey, account, 20, 30, 0.7, base.Add(-24*time.Hour)) // outside range
  1400. startTime := base.Add(-1 * time.Hour)
  1401. endTime := base.Add(2 * time.Hour)
  1402. filters := usagestats.UsageLogFilters{StartTime: &startTime, EndTime: &endTime}
  1403. logs, page, err := s.repo.ListWithFilters(s.ctx, pagination.PaginationParams{Page: 1, PageSize: 10}, filters)
  1404. s.Require().NoError(err, "ListWithFilters time range")
  1405. s.Require().Len(logs, 2)
  1406. s.Require().Equal(int64(2), page.Total)
  1407. }
  1408. func (s *UsageLogRepoSuite) TestListWithFilters_CombinedFilters() {
  1409. user := mustCreateUser(s.T(), s.client, &service.User{Email: "filterscombined@test.com"})
  1410. apiKey := mustCreateApiKey(s.T(), s.client, &service.APIKey{UserID: user.ID, Key: "sk-filterscombined", Name: "k"})
  1411. account := mustCreateAccount(s.T(), s.client, &service.Account{Name: "acc-filterscombined"})
  1412. base := time.Date(2025, 1, 15, 12, 0, 0, 0, time.UTC)
  1413. s.createUsageLog(user, apiKey, account, 10, 20, 0.5, base)
  1414. s.createUsageLog(user, apiKey, account, 15, 25, 0.6, base.Add(1*time.Hour))
  1415. startTime := base.Add(-1 * time.Hour)
  1416. endTime := base.Add(2 * time.Hour)
  1417. filters := usagestats.UsageLogFilters{
  1418. UserID: user.ID,
  1419. APIKeyID: apiKey.ID,
  1420. StartTime: &startTime,
  1421. EndTime: &endTime,
  1422. }
  1423. logs, page, err := s.repo.ListWithFilters(s.ctx, pagination.PaginationParams{Page: 1, PageSize: 10}, filters)
  1424. s.Require().NoError(err, "ListWithFilters combined")
  1425. s.Require().Len(logs, 2)
  1426. s.Require().Equal(int64(2), page.Total)
  1427. }