account_usage_session_window_test.go 3.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135
  1. package service
  2. import (
  3. "context"
  4. "sync"
  5. "testing"
  6. "time"
  7. )
  8. // sessionWindowSyncRepo 记录 syncActiveToPassive 触发的所有写操作。
  9. type sessionWindowSyncRepo struct {
  10. AccountRepository
  11. mu sync.Mutex
  12. extraUpdates []map[string]any
  13. sessionWindowEnds []sessionWindowEndCall
  14. }
  15. type sessionWindowEndCall struct {
  16. AccountID int64
  17. End time.Time
  18. }
  19. func (r *sessionWindowSyncRepo) UpdateExtra(_ context.Context, _ int64, updates map[string]any) error {
  20. r.mu.Lock()
  21. defer r.mu.Unlock()
  22. copied := make(map[string]any, len(updates))
  23. for k, v := range updates {
  24. copied[k] = v
  25. }
  26. r.extraUpdates = append(r.extraUpdates, copied)
  27. return nil
  28. }
  29. func (r *sessionWindowSyncRepo) UpdateSessionWindowEnd(_ context.Context, id int64, end time.Time) error {
  30. r.mu.Lock()
  31. defer r.mu.Unlock()
  32. r.sessionWindowEnds = append(r.sessionWindowEnds, sessionWindowEndCall{AccountID: id, End: end})
  33. return nil
  34. }
  35. func TestEstimateSetupTokenUsage_ExpiredWindowZeroes(t *testing.T) {
  36. t.Parallel()
  37. past := time.Now().Add(-2 * time.Hour)
  38. svc := &AccountUsageService{}
  39. info := svc.estimateSetupTokenUsage(&Account{
  40. SessionWindowEnd: &past,
  41. Extra: map[string]any{
  42. "session_window_utilization": 0.53,
  43. },
  44. })
  45. if info.FiveHour == nil {
  46. t.Fatal("expected non-nil FiveHour info")
  47. }
  48. if info.FiveHour.Utilization != 0 {
  49. t.Fatalf("expected Utilization=0 for expired window, got %v", info.FiveHour.Utilization)
  50. }
  51. if info.FiveHour.ResetsAt != nil {
  52. t.Fatalf("expected ResetsAt=nil for expired window, got %v", info.FiveHour.ResetsAt)
  53. }
  54. if info.FiveHour.RemainingSeconds != 0 {
  55. t.Fatalf("expected RemainingSeconds=0 for expired window, got %v", info.FiveHour.RemainingSeconds)
  56. }
  57. }
  58. func TestEstimateSetupTokenUsage_ActiveWindowPreservesUtilization(t *testing.T) {
  59. t.Parallel()
  60. future := time.Now().Add(3 * time.Hour)
  61. svc := &AccountUsageService{}
  62. info := svc.estimateSetupTokenUsage(&Account{
  63. SessionWindowEnd: &future,
  64. Extra: map[string]any{
  65. "session_window_utilization": 0.53,
  66. },
  67. })
  68. if info.FiveHour == nil {
  69. t.Fatal("expected non-nil FiveHour info")
  70. }
  71. if info.FiveHour.Utilization != 53 {
  72. t.Fatalf("expected Utilization=53, got %v", info.FiveHour.Utilization)
  73. }
  74. if info.FiveHour.ResetsAt == nil || !info.FiveHour.ResetsAt.Equal(future) {
  75. t.Fatalf("expected ResetsAt=%v, got %v", future, info.FiveHour.ResetsAt)
  76. }
  77. if info.FiveHour.RemainingSeconds <= 0 {
  78. t.Fatalf("expected positive RemainingSeconds, got %v", info.FiveHour.RemainingSeconds)
  79. }
  80. }
  81. func TestSyncActiveToPassive_WritesFiveHourSessionWindowEnd(t *testing.T) {
  82. t.Parallel()
  83. repo := &sessionWindowSyncRepo{}
  84. svc := &AccountUsageService{accountRepo: repo}
  85. resetsAt := time.Now().Add(3 * time.Hour).UTC().Truncate(time.Second)
  86. svc.syncActiveToPassive(context.Background(), 42, &UsageInfo{
  87. FiveHour: &UsageProgress{
  88. Utilization: 53,
  89. ResetsAt: &resetsAt,
  90. },
  91. })
  92. repo.mu.Lock()
  93. defer repo.mu.Unlock()
  94. if len(repo.sessionWindowEnds) != 1 {
  95. t.Fatalf("expected 1 UpdateSessionWindowEnd call, got %d", len(repo.sessionWindowEnds))
  96. }
  97. call := repo.sessionWindowEnds[0]
  98. if call.AccountID != 42 {
  99. t.Fatalf("expected AccountID=42, got %d", call.AccountID)
  100. }
  101. if !call.End.Equal(resetsAt) {
  102. t.Fatalf("expected End=%v, got %v", resetsAt, call.End)
  103. }
  104. }
  105. func TestSyncActiveToPassive_SkipsSessionWindowEndWhenResetMissing(t *testing.T) {
  106. t.Parallel()
  107. repo := &sessionWindowSyncRepo{}
  108. svc := &AccountUsageService{accountRepo: repo}
  109. svc.syncActiveToPassive(context.Background(), 99, &UsageInfo{
  110. FiveHour: &UsageProgress{Utilization: 10},
  111. })
  112. repo.mu.Lock()
  113. defer repo.mu.Unlock()
  114. if len(repo.sessionWindowEnds) != 0 {
  115. t.Fatalf("expected no UpdateSessionWindowEnd calls when ResetsAt is nil, got %d", len(repo.sessionWindowEnds))
  116. }
  117. }