tiered_routing_test.go 2.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869
  1. package service
  2. import (
  3. "context"
  4. "testing"
  5. "github.com/stretchr/testify/require"
  6. )
  7. func TestResolveTierByRules(t *testing.T) {
  8. cfg := TieredRoutingSettings{Enabled: true, PremiumUserIDs: []int64{7}, MinBalanceForPremium: 100}
  9. require.Equal(t, TierPremium, resolveTierByRules(&User{ID: 7}, cfg))
  10. require.Equal(t, TierPremium, resolveTierByRules(&User{ID: 8, Balance: 150}, cfg))
  11. require.Equal(t, TierStandard, resolveTierByRules(&User{ID: 8, Balance: 50}, cfg))
  12. require.Equal(t, TierStandard, resolveTierByRules(nil, cfg))
  13. }
  14. func TestTieredRoutingDisabled(t *testing.T) {
  15. svc := NewTieredRoutingService(nil)
  16. require.Equal(t, TierStandard, svc.ResolveUserTier(&User{ID: 7}))
  17. require.Equal(t, TierStandard, svc.ResolvePool(&User{ID: 7, Balance: 9999}))
  18. }
  19. func TestAccountPool(t *testing.T) {
  20. require.Equal(t, TierStandard, AccountPool(nil))
  21. require.Equal(t, TierStandard, AccountPool(&Account{}))
  22. require.Equal(t, TierPremium, AccountPool(&Account{Extra: map[string]any{"pool": "premium"}}))
  23. require.Equal(t, TierPremium, AccountPool(&Account{Extra: map[string]any{"pool": " Premium "}}))
  24. require.Equal(t, TierStandard, AccountPool(&Account{Extra: map[string]any{"pool": "gold"}}))
  25. }
  26. func TestFilterTierPoolAccounts(t *testing.T) {
  27. accounts := []Account{
  28. {ID: 1, Extra: map[string]any{"pool": "premium"}},
  29. {ID: 2},
  30. }
  31. // standard 不过滤
  32. require.Len(t, FilterTierPoolAccounts(accounts, TierStandard), 2)
  33. // premium 命中池
  34. got := FilterTierPoolAccounts(accounts, TierPremium)
  35. require.Len(t, got, 1)
  36. require.Equal(t, int64(1), got[0].ID)
  37. // 池空时回退全量
  38. require.Len(t, FilterTierPoolAccounts(accounts[1:], TierPremium), 1)
  39. // 空输入
  40. require.Empty(t, FilterTierPoolAccounts(nil, TierPremium))
  41. }
  42. func TestTierPoolContext(t *testing.T) {
  43. require.Equal(t, TierStandard, TierPoolFromContext(nil))
  44. require.Equal(t, TierStandard, TierPoolFromContext(context.Background()))
  45. ctx := WithTierPool(context.Background(), TierPremium)
  46. require.Equal(t, TierPremium, TierPoolFromContext(ctx))
  47. }
  48. func TestTieredRoutingSettingsRoundtrip(t *testing.T) {
  49. repo := &stubHealthSettingRepo{}
  50. svc := NewTieredRoutingService(repo)
  51. got := svc.GetSettings(context.Background())
  52. require.False(t, got.Enabled)
  53. updated, err := svc.UpdateSettings(context.Background(), TieredRoutingSettings{
  54. Enabled: true, PremiumUserIDs: []int64{1, 1, -5}, MinBalanceForPremium: -1,
  55. })
  56. require.NoError(t, err)
  57. require.True(t, updated.Enabled)
  58. require.Equal(t, []int64{1}, updated.PremiumUserIDs)
  59. require.Equal(t, 0.0, updated.MinBalanceForPremium)
  60. }