admin_service_clear_error_test.go 2.6 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788
  1. //go:build unit
  2. package service
  3. import (
  4. "context"
  5. "testing"
  6. "time"
  7. "github.com/stretchr/testify/require"
  8. )
  9. type accountRepoStubForClearAccountError struct {
  10. mockAccountRepoForGemini
  11. account *Account
  12. clearErrorCalls int
  13. clearRateLimitCalls int
  14. clearAntigravityCalls int
  15. clearModelRateLimitCalls int
  16. clearTempUnschedCalls int
  17. }
  18. func (r *accountRepoStubForClearAccountError) GetByID(ctx context.Context, id int64) (*Account, error) {
  19. return r.account, nil
  20. }
  21. func (r *accountRepoStubForClearAccountError) ClearError(ctx context.Context, id int64) error {
  22. r.clearErrorCalls++
  23. r.account.Status = StatusActive
  24. r.account.ErrorMessage = ""
  25. return nil
  26. }
  27. func (r *accountRepoStubForClearAccountError) ClearRateLimit(ctx context.Context, id int64) error {
  28. r.clearRateLimitCalls++
  29. r.account.RateLimitedAt = nil
  30. r.account.RateLimitResetAt = nil
  31. return nil
  32. }
  33. func (r *accountRepoStubForClearAccountError) ClearAntigravityQuotaScopes(ctx context.Context, id int64) error {
  34. r.clearAntigravityCalls++
  35. return nil
  36. }
  37. func (r *accountRepoStubForClearAccountError) ClearModelRateLimits(ctx context.Context, id int64) error {
  38. r.clearModelRateLimitCalls++
  39. return nil
  40. }
  41. func (r *accountRepoStubForClearAccountError) ClearTempUnschedulable(ctx context.Context, id int64) error {
  42. r.clearTempUnschedCalls++
  43. r.account.TempUnschedulableUntil = nil
  44. r.account.TempUnschedulableReason = ""
  45. return nil
  46. }
  47. func TestAdminService_ClearAccountError_AlsoClearsRecoverableRuntimeState(t *testing.T) {
  48. until := time.Now().Add(10 * time.Minute)
  49. resetAt := time.Now().Add(5 * time.Minute)
  50. repo := &accountRepoStubForClearAccountError{
  51. account: &Account{
  52. ID: 31,
  53. Platform: PlatformOpenAI,
  54. Type: AccountTypeOAuth,
  55. Status: StatusError,
  56. ErrorMessage: "refresh failed",
  57. RateLimitResetAt: &resetAt,
  58. TempUnschedulableUntil: &until,
  59. TempUnschedulableReason: "missing refresh token",
  60. },
  61. }
  62. blocker := &runtimeBlockRecorder{}
  63. svc := &adminServiceImpl{accountRepo: repo, runtimeBlocker: blocker}
  64. updated, err := svc.ClearAccountError(context.Background(), 31)
  65. require.NoError(t, err)
  66. require.NotNil(t, updated)
  67. require.Equal(t, 1, repo.clearErrorCalls)
  68. require.Equal(t, 1, repo.clearRateLimitCalls)
  69. require.Equal(t, 1, repo.clearAntigravityCalls)
  70. require.Equal(t, 1, repo.clearModelRateLimitCalls)
  71. require.Equal(t, 1, repo.clearTempUnschedCalls)
  72. require.Nil(t, updated.RateLimitResetAt)
  73. require.Nil(t, updated.TempUnschedulableUntil)
  74. require.Empty(t, updated.TempUnschedulableReason)
  75. require.Equal(t, []int64{31}, blocker.clearedIDs)
  76. }