gateway_model_availability_test.go 9.1 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251
  1. //go:build unit
  2. package service
  3. import (
  4. "context"
  5. "testing"
  6. "time"
  7. "github.com/stretchr/testify/require"
  8. )
  9. func TestDiagnoseModelAvailabilityForPlatform_NoModel_AlwaysAvailable(t *testing.T) {
  10. repo := &mockAccountRepoForPlatform{accounts: nil, accountsByID: map[int64]*Account{}}
  11. svc := &GatewayService{accountRepo: repo, cfg: testConfig()}
  12. diag := svc.DiagnoseModelAvailabilityForPlatform(context.Background(), nil, "", PlatformOpenAI)
  13. require.True(t, diag.HasAccountsInPool, "empty model must return HasAccountsInPool=true so caller stays on 503")
  14. require.True(t, diag.HasModelSupport, "empty model must return HasModelSupport=true so caller stays on 503")
  15. }
  16. func TestDiagnoseModelAvailabilityForPlatform_EmptyPlatform_AlwaysAvailable(t *testing.T) {
  17. repo := &mockAccountRepoForPlatform{accounts: nil, accountsByID: map[int64]*Account{}}
  18. svc := &GatewayService{accountRepo: repo, cfg: testConfig()}
  19. diag := svc.DiagnoseModelAvailabilityForPlatform(context.Background(), nil, "gpt-5", "")
  20. require.True(t, diag.HasAccountsInPool)
  21. require.True(t, diag.HasModelSupport, "empty platform must fall back to {true,true} so caller stays on 503")
  22. }
  23. func TestDiagnoseModelAvailabilityForPlatform_NilReceiver(t *testing.T) {
  24. var svc *GatewayService
  25. diag := svc.DiagnoseModelAvailabilityForPlatform(context.Background(), nil, "gpt-5", PlatformOpenAI)
  26. require.True(t, diag.HasAccountsInPool)
  27. require.True(t, diag.HasModelSupport)
  28. }
  29. func TestDiagnoseModelAvailabilityForPlatform_NoAccountsInPool(t *testing.T) {
  30. repo := &mockAccountRepoForPlatform{accounts: nil, accountsByID: map[int64]*Account{}}
  31. svc := &GatewayService{accountRepo: repo, cfg: testConfig()}
  32. diag := svc.DiagnoseModelAvailabilityForPlatform(context.Background(), nil, "gpt-5", PlatformOpenAI)
  33. require.False(t, diag.HasAccountsInPool)
  34. require.False(t, diag.HasModelSupport, "no accounts means no support; caller stays on 503 (empty-pool branch)")
  35. }
  36. func TestDiagnoseModelAvailabilityForPlatform_ExplicitMappingMatches(t *testing.T) {
  37. repo := &mockAccountRepoForPlatform{
  38. accounts: []Account{
  39. {
  40. ID: 1,
  41. Platform: PlatformOpenAI,
  42. Status: StatusActive,
  43. Schedulable: true,
  44. Credentials: map[string]any{
  45. "model_mapping": map[string]any{"gpt-5.1-codex-mini": "gpt-5.1-codex-mini"},
  46. },
  47. },
  48. },
  49. accountsByID: map[int64]*Account{},
  50. }
  51. for i := range repo.accounts {
  52. repo.accountsByID[repo.accounts[i].ID] = &repo.accounts[i]
  53. }
  54. svc := &GatewayService{accountRepo: repo, cfg: testConfig()}
  55. diag := svc.DiagnoseModelAvailabilityForPlatform(context.Background(), nil, "gpt-5.1-codex-mini", PlatformOpenAI)
  56. require.True(t, diag.HasAccountsInPool)
  57. require.True(t, diag.HasModelSupport)
  58. }
  59. func TestDiagnoseModelAvailabilityForPlatform_EmptyMappingAllowsAll(t *testing.T) {
  60. repo := &mockAccountRepoForPlatform{
  61. accounts: []Account{
  62. {ID: 1, Platform: PlatformOpenAI, Status: StatusActive, Schedulable: true /* no ModelMapping = allow all */},
  63. },
  64. accountsByID: map[int64]*Account{},
  65. }
  66. for i := range repo.accounts {
  67. repo.accountsByID[repo.accounts[i].ID] = &repo.accounts[i]
  68. }
  69. svc := &GatewayService{accountRepo: repo, cfg: testConfig()}
  70. diag := svc.DiagnoseModelAvailabilityForPlatform(context.Background(), nil, "gpt-5.1-codex-mini", PlatformOpenAI)
  71. require.True(t, diag.HasModelSupport, "empty model_mapping must be treated as 'allow all' (Account.IsModelSupported semantics)")
  72. }
  73. func TestDiagnoseModelAvailabilityForPlatform_WildcardMappingMatches(t *testing.T) {
  74. repo := &mockAccountRepoForPlatform{
  75. accounts: []Account{
  76. {
  77. ID: 1,
  78. Platform: PlatformOpenAI,
  79. Status: StatusActive,
  80. Schedulable: true,
  81. Credentials: map[string]any{
  82. "model_mapping": map[string]any{"*": "gpt-5"},
  83. },
  84. },
  85. },
  86. accountsByID: map[int64]*Account{},
  87. }
  88. for i := range repo.accounts {
  89. repo.accountsByID[repo.accounts[i].ID] = &repo.accounts[i]
  90. }
  91. svc := &GatewayService{accountRepo: repo, cfg: testConfig()}
  92. diag := svc.DiagnoseModelAvailabilityForPlatform(context.Background(), nil, "gpt-5.1-codex-mini", PlatformOpenAI)
  93. require.True(t, diag.HasModelSupport, "wildcard mapping must classify the request as 'serviceable'")
  94. }
  95. func TestDiagnoseModelAvailabilityForPlatform_NoMatchingModel_ReturnsNotFoundSignal(t *testing.T) {
  96. groupID := int64(42)
  97. repo := &mockAccountRepoForPlatform{
  98. accounts: []Account{
  99. {
  100. ID: 1,
  101. Platform: PlatformOpenAI,
  102. Status: StatusActive,
  103. Schedulable: true,
  104. AccountGroups: []AccountGroup{
  105. {GroupID: groupID},
  106. },
  107. Credentials: map[string]any{"model_mapping": map[string]any{"gpt-5": "gpt-5"}},
  108. },
  109. {
  110. ID: 2,
  111. Platform: PlatformOpenAI,
  112. Status: StatusActive,
  113. Schedulable: true,
  114. AccountGroups: []AccountGroup{
  115. {GroupID: groupID},
  116. },
  117. Credentials: map[string]any{"model_mapping": map[string]any{"gpt-5-mini": "gpt-5-mini"}},
  118. },
  119. },
  120. accountsByID: map[int64]*Account{},
  121. }
  122. for i := range repo.accounts {
  123. repo.accountsByID[repo.accounts[i].ID] = &repo.accounts[i]
  124. }
  125. svc := &GatewayService{accountRepo: repo, cfg: testConfig()}
  126. diag := svc.DiagnoseModelAvailabilityForPlatform(context.Background(), &groupID, "gpt-5.1-codex-mini", PlatformOpenAI)
  127. require.True(t, diag.HasAccountsInPool, "group has OpenAI accounts")
  128. require.False(t, diag.HasModelSupport, "no account mapping admits the requested model — handler should return 404")
  129. }
  130. func TestDiagnoseModelAvailabilityForPlatform_RateLimitedSupportingAccountRemainsConfigured(t *testing.T) {
  131. groupID := int64(42)
  132. cooldownUntil := time.Now().Add(time.Hour)
  133. repo := &mockAccountRepoForPlatform{
  134. accounts: []Account{
  135. {
  136. ID: 1,
  137. Platform: PlatformAnthropic,
  138. Status: StatusActive,
  139. Schedulable: true,
  140. RateLimitResetAt: &cooldownUntil,
  141. OverloadUntil: &cooldownUntil,
  142. TempUnschedulableUntil: &cooldownUntil,
  143. AccountGroups: []AccountGroup{{GroupID: groupID}},
  144. Credentials: map[string]any{
  145. "model_mapping": map[string]any{"claude-opus-4-8": "claude-opus-4-8"},
  146. },
  147. },
  148. },
  149. accountsByID: map[int64]*Account{},
  150. }
  151. require.False(t, repo.accounts[0].IsSchedulable(), "test account must be excluded from normal scheduling while cooling down")
  152. svc := &GatewayService{
  153. accountRepo: repo,
  154. cfg: testConfig(),
  155. schedulerSnapshot: &SchedulerSnapshotService{}, // diagnosis must bypass the transient-only snapshot
  156. }
  157. diag := svc.DiagnoseModelAvailabilityForPlatform(context.Background(), &groupID, "claude-opus-4-8", PlatformAnthropic)
  158. require.True(t, diag.HasAccountsInPool)
  159. require.True(t, diag.HasModelSupport, "a configured model remains supported while every matching account is temporarily cooling down")
  160. }
  161. func TestOpenAIDiagnoseModelAvailabilityForPlatform_RateLimitedSupportingAccountRemainsConfigured(t *testing.T) {
  162. groupID := int64(43)
  163. cooldownUntil := time.Now().Add(time.Hour)
  164. repo := &mockAccountRepoForPlatform{
  165. accounts: []Account{
  166. {
  167. ID: 2,
  168. Platform: PlatformOpenAI,
  169. Status: StatusActive,
  170. Schedulable: true,
  171. RateLimitResetAt: &cooldownUntil,
  172. OverloadUntil: &cooldownUntil,
  173. TempUnschedulableUntil: &cooldownUntil,
  174. AccountGroups: []AccountGroup{{GroupID: groupID}},
  175. Credentials: map[string]any{
  176. "model_mapping": map[string]any{"claude-opus-4-8": "claude-opus-4-8"},
  177. },
  178. },
  179. },
  180. accountsByID: map[int64]*Account{},
  181. }
  182. require.False(t, repo.accounts[0].IsSchedulable(), "test account must be excluded from normal scheduling while cooling down")
  183. svc := &OpenAIGatewayService{
  184. accountRepo: repo,
  185. cfg: testConfig(),
  186. schedulerSnapshot: &SchedulerSnapshotService{}, // diagnosis must bypass the transient-only snapshot
  187. }
  188. diag := svc.DiagnoseModelAvailabilityForPlatform(context.Background(), &groupID, "claude-opus-4-8", PlatformOpenAI)
  189. require.True(t, diag.HasAccountsInPool)
  190. require.True(t, diag.HasModelSupport, "OpenAI-compatible diagnosis must keep transiently limited supporting accounts in the configured pool")
  191. }
  192. func TestDiagnoseModelAvailabilityForPlatform_WrongPlatformFiltersOut(t *testing.T) {
  193. // Group has only Anthropic accounts; user routes to OpenAI gateway.
  194. // Diagnosis must NOT see Anthropic accounts (listSchedulableAccounts filters
  195. // by platform), so HasAccountsInPool is false and the caller stays on 503.
  196. repo := &mockAccountRepoForPlatform{
  197. accounts: []Account{
  198. {
  199. ID: 1,
  200. Platform: PlatformAnthropic,
  201. Status: StatusActive,
  202. Schedulable: true,
  203. Credentials: map[string]any{"model_mapping": map[string]any{"claude-sonnet-4-5": "claude-sonnet-4-5"}},
  204. },
  205. },
  206. accountsByID: map[int64]*Account{},
  207. }
  208. for i := range repo.accounts {
  209. repo.accountsByID[repo.accounts[i].ID] = &repo.accounts[i]
  210. }
  211. svc := &GatewayService{accountRepo: repo, cfg: testConfig()}
  212. diag := svc.DiagnoseModelAvailabilityForPlatform(context.Background(), nil, "gpt-5", PlatformOpenAI)
  213. require.False(t, diag.HasAccountsInPool, "OpenAI route must not see Anthropic accounts in pool")
  214. require.False(t, diag.HasModelSupport)
  215. }