scheduler_cache_test.go 3.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110
  1. package repository
  2. import (
  3. "strings"
  4. "testing"
  5. "github.com/Wei-Shaw/sub2api/internal/service"
  6. "github.com/stretchr/testify/require"
  7. )
  8. func TestFilterSchedulerCredentialsKeepsSubscriptionPlanType(t *testing.T) {
  9. filtered := filterSchedulerCredentials(map[string]any{
  10. "plan_type": "plus",
  11. "access_token": "secret-access-token",
  12. "refresh_token": "secret-refresh-token",
  13. })
  14. require.Equal(t, "plus", filtered["plan_type"])
  15. require.NotContains(t, filtered, "access_token")
  16. require.NotContains(t, filtered, "refresh_token")
  17. }
  18. func TestSchedulerMetadataAccountKeepsOpenAISubscriptionIdentity(t *testing.T) {
  19. account := service.Account{
  20. ID: 24,
  21. Platform: service.PlatformOpenAI,
  22. Type: service.AccountTypeOAuth,
  23. Credentials: map[string]any{
  24. "plan_type": "plus",
  25. "access_token": "secret-access-token",
  26. },
  27. }
  28. metadata := buildSchedulerMetadataAccount(account)
  29. require.True(t, metadata.IsOpenAIChatGPTSubscription())
  30. require.Empty(t, metadata.GetCredential("access_token"))
  31. }
  32. func TestSchedulerMetadataAccountProjectsUpstreamBillingProbe(t *testing.T) {
  33. lastError := strings.Repeat("upstream diagnostic ", 512)
  34. probe := map[string]any{
  35. "status": "ok",
  36. "data": map[string]any{
  37. "billing_scope": "token",
  38. "resolved_rate_multiplier": 0.03,
  39. "peak_rate_enabled": true,
  40. "peak_start": "09:00",
  41. "peak_end": "18:00",
  42. "peak_rate_multiplier": 2.0,
  43. "timezone": "Asia/Shanghai",
  44. "effective_rate_multiplier": 0.03,
  45. "remote_diagnostic": lastError,
  46. },
  47. "received_at": "2026-07-13T10:00:00Z",
  48. "fresh_until": "2026-07-13T11:00:00Z",
  49. "next_probe_at": "2026-07-13T10:30:00Z",
  50. "http_status": 502,
  51. "last_error": lastError,
  52. }
  53. account := service.Account{
  54. ID: 42,
  55. Extra: map[string]any{
  56. "upstream_billing_probe": probe,
  57. "unused_large_field": "drop-me",
  58. },
  59. }
  60. metadata := buildSchedulerMetadataAccount(account)
  61. fullPayload, metaPayload, err := marshalSchedulerCacheAccount(account)
  62. require.NoError(t, err)
  63. filtered, ok := metadata.Extra["upstream_billing_probe"].(map[string]any)
  64. require.True(t, ok)
  65. require.Equal(t, "ok", filtered["status"])
  66. require.Equal(t, "2026-07-13T10:00:00Z", filtered["received_at"])
  67. require.Equal(t, "2026-07-13T11:00:00Z", filtered["fresh_until"])
  68. require.Equal(t, "2026-07-13T10:30:00Z", filtered["next_probe_at"])
  69. require.NotContains(t, filtered, "http_status")
  70. require.NotContains(t, filtered, "last_error")
  71. filteredData, ok := filtered["data"].(map[string]any)
  72. require.True(t, ok)
  73. require.Equal(t, "token", filteredData["billing_scope"])
  74. require.Equal(t, 0.03, filteredData["resolved_rate_multiplier"])
  75. require.Equal(t, true, filteredData["peak_rate_enabled"])
  76. require.Equal(t, "09:00", filteredData["peak_start"])
  77. require.Equal(t, "18:00", filteredData["peak_end"])
  78. require.Equal(t, 2.0, filteredData["peak_rate_multiplier"])
  79. require.Equal(t, "Asia/Shanghai", filteredData["timezone"])
  80. require.NotContains(t, filteredData, "effective_rate_multiplier")
  81. require.NotContains(t, filteredData, "remote_diagnostic")
  82. require.NotContains(t, metadata.Extra, "unused_large_field")
  83. require.Contains(t, string(fullPayload), lastError)
  84. require.NotContains(t, string(metaPayload), "last_error")
  85. require.Less(t, len(metaPayload)*4, len(fullPayload))
  86. }
  87. func TestSchedulerMetadataAccountDropsInvalidUpstreamBillingProbe(t *testing.T) {
  88. for _, probe := range []any{
  89. "invalid",
  90. map[string]any{},
  91. map[string]any{"status": ""},
  92. } {
  93. metadata := buildSchedulerMetadataAccount(service.Account{
  94. Extra: map[string]any{service.UpstreamBillingProbeExtraKey: probe},
  95. })
  96. require.NotContains(t, metadata.Extra, service.UpstreamBillingProbeExtraKey)
  97. }
  98. }