user_platform_quota_upsert_test.go 4.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148
  1. //go:build integration
  2. package repository
  3. import (
  4. "context"
  5. "testing"
  6. "time"
  7. "github.com/Wei-Shaw/sub2api/ent/userplatformquota"
  8. "github.com/stretchr/testify/require"
  9. )
  10. func TestUpsertForUser_NewUserInsertsAllRecords(t *testing.T) {
  11. ctx := context.Background()
  12. client := testEntClient(t)
  13. userID := mustCreateUserForQuota(t, client)
  14. repo := NewUserPlatformQuotaRepository(client)
  15. daily := 10.0
  16. weekly := 50.0
  17. monthly := 200.0
  18. records := []UserPlatformQuotaRecord{
  19. {UserID: userID, Platform: "anthropic", DailyLimitUSD: &daily, WeeklyLimitUSD: &weekly, MonthlyLimitUSD: &monthly},
  20. {UserID: userID, Platform: "openai", DailyLimitUSD: &daily},
  21. }
  22. require.NoError(t, repo.UpsertForUser(ctx, userID, records))
  23. got, err := repo.ListByUser(ctx, userID)
  24. require.NoError(t, err)
  25. require.Len(t, got, 2)
  26. }
  27. func TestUpsertForUser_PartialUpdateSoftDeletesMissingPlatforms(t *testing.T) {
  28. ctx := context.Background()
  29. client := testEntClient(t)
  30. userID := mustCreateUserForQuota(t, client)
  31. repo := NewUserPlatformQuotaRepository(client)
  32. d1 := 10.0
  33. d2 := 20.0
  34. require.NoError(t, repo.UpsertForUser(ctx, userID, []UserPlatformQuotaRecord{
  35. {UserID: userID, Platform: "anthropic", DailyLimitUSD: &d1},
  36. {UserID: userID, Platform: "openai", DailyLimitUSD: &d1},
  37. }))
  38. require.NoError(t, repo.UpsertForUser(ctx, userID, []UserPlatformQuotaRecord{
  39. {UserID: userID, Platform: "anthropic", DailyLimitUSD: &d2},
  40. {UserID: userID, Platform: "gemini", DailyLimitUSD: &d1},
  41. }))
  42. active, err := repo.ListByUser(ctx, userID)
  43. require.NoError(t, err)
  44. platforms := map[string]float64{}
  45. for _, r := range active {
  46. require.NotNil(t, r.DailyLimitUSD)
  47. platforms[r.Platform] = *r.DailyLimitUSD
  48. }
  49. require.Len(t, platforms, 2)
  50. require.InDelta(t, 20.0, platforms["anthropic"], 1e-9)
  51. require.InDelta(t, 10.0, platforms["gemini"], 1e-9)
  52. _, openaiActive := platforms["openai"]
  53. require.False(t, openaiActive, "openai should be soft-deleted")
  54. }
  55. func TestUpsertForUser_PreservesUsageAndWindowStart(t *testing.T) {
  56. ctx := context.Background()
  57. client := testEntClient(t)
  58. userID := mustCreateUserForQuota(t, client)
  59. repo := NewUserPlatformQuotaRepository(client)
  60. d := 10.0
  61. require.NoError(t, repo.UpsertForUser(ctx, userID, []UserPlatformQuotaRecord{
  62. {UserID: userID, Platform: "anthropic", DailyLimitUSD: &d},
  63. }))
  64. now := time.Date(2026, 5, 22, 10, 0, 0, 0, time.UTC)
  65. require.NoError(t, repo.IncrementUsageWithReset(ctx, userID, "anthropic", 3.5, now))
  66. newD := 50.0
  67. require.NoError(t, repo.UpsertForUser(ctx, userID, []UserPlatformQuotaRecord{
  68. {UserID: userID, Platform: "anthropic", DailyLimitUSD: &newD},
  69. }))
  70. rec, err := repo.GetByUserPlatform(ctx, userID, "anthropic")
  71. require.NoError(t, err)
  72. require.NotNil(t, rec)
  73. require.InDelta(t, 50.0, *rec.DailyLimitUSD, 1e-9, "limit should update")
  74. require.InDelta(t, 3.5, rec.DailyUsageUSD, 1e-9, "usage must be preserved")
  75. require.NotNil(t, rec.DailyWindowStart, "window_start must be preserved")
  76. }
  77. func TestUpsertForUser_ReactivatesSoftDeleted(t *testing.T) {
  78. ctx := context.Background()
  79. client := testEntClient(t)
  80. userID := mustCreateUserForQuota(t, client)
  81. repo := NewUserPlatformQuotaRepository(client)
  82. d := 10.0
  83. require.NoError(t, repo.UpsertForUser(ctx, userID, []UserPlatformQuotaRecord{
  84. {UserID: userID, Platform: "anthropic", DailyLimitUSD: &d},
  85. }))
  86. require.NoError(t, repo.UpsertForUser(ctx, userID, []UserPlatformQuotaRecord{}))
  87. gone, err := repo.GetByUserPlatform(ctx, userID, "anthropic")
  88. require.NoError(t, err)
  89. require.Nil(t, gone, "anthropic should be soft-deleted (not active)")
  90. d2 := 20.0
  91. require.NoError(t, repo.UpsertForUser(ctx, userID, []UserPlatformQuotaRecord{
  92. {UserID: userID, Platform: "anthropic", DailyLimitUSD: &d2},
  93. }))
  94. back, err := repo.GetByUserPlatform(ctx, userID, "anthropic")
  95. require.NoError(t, err)
  96. require.NotNil(t, back, "anthropic should be active again")
  97. require.InDelta(t, 20.0, *back.DailyLimitUSD, 1e-9)
  98. allRows, err := client.UserPlatformQuota.Query().
  99. Where(userplatformquota.UserIDEQ(userID), userplatformquota.PlatformEQ("anthropic")).
  100. All(ctx)
  101. require.NoError(t, err)
  102. activeCount := 0
  103. for _, r := range allRows {
  104. if r.DeletedAt == nil {
  105. activeCount++
  106. }
  107. }
  108. require.Equal(t, 1, activeCount, "should have exactly one active row")
  109. }
  110. func TestUpsertForUser_EmptyClearsAll(t *testing.T) {
  111. ctx := context.Background()
  112. client := testEntClient(t)
  113. userID := mustCreateUserForQuota(t, client)
  114. repo := NewUserPlatformQuotaRepository(client)
  115. d := 10.0
  116. require.NoError(t, repo.UpsertForUser(ctx, userID, []UserPlatformQuotaRecord{
  117. {UserID: userID, Platform: "anthropic", DailyLimitUSD: &d},
  118. {UserID: userID, Platform: "openai", DailyLimitUSD: &d},
  119. }))
  120. require.NoError(t, repo.UpsertForUser(ctx, userID, []UserPlatformQuotaRecord{}))
  121. got, err := repo.ListByUser(ctx, userID)
  122. require.NoError(t, err)
  123. require.Empty(t, got)
  124. }