gateway_channel_restriction_fallback_test.go 3.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130
  1. //go:build unit
  2. package service
  3. import (
  4. "context"
  5. "testing"
  6. "github.com/Wei-Shaw/sub2api/internal/pkg/ctxkey"
  7. "github.com/stretchr/testify/require"
  8. )
  9. func TestSelectAccountForModelWithExclusions_UsesFallbackGroupForChannelRestriction(t *testing.T) {
  10. t.Parallel()
  11. groupID := int64(10)
  12. fallbackID := int64(11)
  13. ch := Channel{
  14. ID: 1,
  15. Status: StatusActive,
  16. GroupIDs: []int64{fallbackID},
  17. RestrictModels: true,
  18. ModelPricing: []ChannelModelPricing{
  19. {Platform: PlatformAnthropic, Models: []string{"claude-sonnet-4-6"}},
  20. },
  21. }
  22. channelSvc := newTestChannelService(makeStandardRepo(ch, map[int64]string{
  23. fallbackID: PlatformAnthropic,
  24. }))
  25. accountRepo := &mockAccountRepoForPlatform{
  26. accounts: []Account{
  27. {ID: 1, Platform: PlatformAnthropic, Priority: 1, Status: StatusActive, Schedulable: true},
  28. },
  29. accountsByID: map[int64]*Account{},
  30. }
  31. for i := range accountRepo.accounts {
  32. accountRepo.accountsByID[accountRepo.accounts[i].ID] = &accountRepo.accounts[i]
  33. }
  34. groupRepo := &mockGroupRepoForGateway{
  35. groups: map[int64]*Group{
  36. groupID: {
  37. ID: groupID,
  38. Platform: PlatformAnthropic,
  39. Status: StatusActive,
  40. ClaudeCodeOnly: true,
  41. FallbackGroupID: &fallbackID,
  42. Hydrated: true,
  43. },
  44. fallbackID: {
  45. ID: fallbackID,
  46. Platform: PlatformAnthropic,
  47. Status: StatusActive,
  48. Hydrated: true,
  49. },
  50. },
  51. }
  52. svc := &GatewayService{
  53. accountRepo: accountRepo,
  54. groupRepo: groupRepo,
  55. channelService: channelSvc,
  56. cfg: testConfig(),
  57. }
  58. ctx := context.WithValue(context.Background(), ctxkey.Group, groupRepo.groups[groupID])
  59. account, err := svc.SelectAccountForModelWithExclusions(ctx, &groupID, "", "claude-sonnet-4-6", nil)
  60. require.NoError(t, err)
  61. require.NotNil(t, account)
  62. require.Equal(t, int64(1), account.ID)
  63. }
  64. func TestSelectAccountWithLoadAwareness_UsesFallbackGroupForChannelRestriction(t *testing.T) {
  65. t.Parallel()
  66. groupID := int64(10)
  67. fallbackID := int64(11)
  68. ch := Channel{
  69. ID: 1,
  70. Status: StatusActive,
  71. GroupIDs: []int64{fallbackID},
  72. RestrictModels: true,
  73. ModelPricing: []ChannelModelPricing{
  74. {Platform: PlatformAnthropic, Models: []string{"claude-sonnet-4-6"}},
  75. },
  76. }
  77. channelSvc := newTestChannelService(makeStandardRepo(ch, map[int64]string{
  78. fallbackID: PlatformAnthropic,
  79. }))
  80. accountRepo := &mockAccountRepoForPlatform{
  81. accounts: []Account{
  82. {ID: 1, Platform: PlatformAnthropic, Priority: 1, Status: StatusActive, Schedulable: true},
  83. },
  84. accountsByID: map[int64]*Account{},
  85. }
  86. for i := range accountRepo.accounts {
  87. accountRepo.accountsByID[accountRepo.accounts[i].ID] = &accountRepo.accounts[i]
  88. }
  89. groupRepo := &mockGroupRepoForGateway{
  90. groups: map[int64]*Group{
  91. groupID: {
  92. ID: groupID,
  93. Platform: PlatformAnthropic,
  94. Status: StatusActive,
  95. ClaudeCodeOnly: true,
  96. FallbackGroupID: &fallbackID,
  97. Hydrated: true,
  98. },
  99. fallbackID: {
  100. ID: fallbackID,
  101. Platform: PlatformAnthropic,
  102. Status: StatusActive,
  103. Hydrated: true,
  104. },
  105. },
  106. }
  107. svc := &GatewayService{
  108. accountRepo: accountRepo,
  109. groupRepo: groupRepo,
  110. channelService: channelSvc,
  111. cfg: testConfig(),
  112. }
  113. ctx := context.WithValue(context.Background(), ctxkey.Group, groupRepo.groups[groupID])
  114. result, err := svc.SelectAccountWithLoadAwareness(ctx, &groupID, "", "claude-sonnet-4-6", nil, "", 0)
  115. require.NoError(t, err)
  116. require.NotNil(t, result)
  117. require.NotNil(t, result.Account)
  118. require.Equal(t, int64(1), result.Account.ID)
  119. }