simple_mode_default_groups_integration_test.go 5.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178
  1. //go:build integration
  2. package repository
  3. import (
  4. "context"
  5. "testing"
  6. "time"
  7. "github.com/Wei-Shaw/sub2api/ent/group"
  8. "github.com/Wei-Shaw/sub2api/internal/service"
  9. "github.com/stretchr/testify/require"
  10. )
  11. func TestEnsureSimpleModeDefaultGroups_CreatesMissingDefaults(t *testing.T) {
  12. ctx := context.Background()
  13. tx := testEntTx(t)
  14. client := tx.Client()
  15. seedCtx, cancel := context.WithTimeout(ctx, 10*time.Second)
  16. defer cancel()
  17. require.NoError(t, ensureSimpleModeDefaultGroups(seedCtx, client))
  18. assertGroupExists := func(name string) {
  19. exists, err := client.Group.Query().Where(group.NameEQ(name), group.DeletedAtIsNil()).Exist(seedCtx)
  20. require.NoError(t, err)
  21. require.True(t, exists, "expected group %s to exist", name)
  22. }
  23. assertGroupExists(service.PlatformAnthropic + "-default")
  24. assertGroupExists(service.PlatformOpenAI + "-default")
  25. assertGroupExists(service.PlatformGemini + "-default")
  26. assertGroupExists(service.PlatformAntigravity + "-default-1")
  27. assertGroupExists(service.PlatformAntigravity + "-default-2")
  28. grokDefault, err := client.Group.Query().
  29. Where(group.NameEQ(service.PlatformGrok+"-default"), group.DeletedAtIsNil()).
  30. Only(seedCtx)
  31. require.NoError(t, err)
  32. require.True(t, grokDefault.AllowImageGeneration)
  33. }
  34. func TestEnsureSimpleModeDefaultGroups_BackfillsOnlyAutoCreatedGrokDefault(t *testing.T) {
  35. ctx := context.Background()
  36. tx := testEntTx(t)
  37. client := tx.Client()
  38. seedCtx, cancel := context.WithTimeout(ctx, 10*time.Second)
  39. defer cancel()
  40. autoDefault, err := client.Group.Create().
  41. SetName(service.PlatformGrok + "-default").
  42. SetDescription("Auto-created default group").
  43. SetPlatform(service.PlatformGrok).
  44. SetStatus(service.StatusActive).
  45. SetSubscriptionType(service.SubscriptionTypeStandard).
  46. SetRateMultiplier(1.0).
  47. SetIsExclusive(false).
  48. SetAllowImageGeneration(false).
  49. Save(seedCtx)
  50. require.NoError(t, err)
  51. operatorGroup, err := client.Group.Create().
  52. SetName("operator-grok-images-disabled-" + time.Now().Format(time.RFC3339Nano)).
  53. SetDescription("Operator-managed group").
  54. SetPlatform(service.PlatformGrok).
  55. SetStatus(service.StatusActive).
  56. SetSubscriptionType(service.SubscriptionTypeStandard).
  57. SetRateMultiplier(1.0).
  58. SetIsExclusive(false).
  59. SetAllowImageGeneration(false).
  60. Save(seedCtx)
  61. require.NoError(t, err)
  62. require.NoError(t, ensureSimpleModeDefaultGroups(seedCtx, client))
  63. autoDefault, err = client.Group.Get(seedCtx, autoDefault.ID)
  64. require.NoError(t, err)
  65. require.True(t, autoDefault.AllowImageGeneration)
  66. operatorGroup, err = client.Group.Get(seedCtx, operatorGroup.ID)
  67. require.NoError(t, err)
  68. require.False(t, operatorGroup.AllowImageGeneration, "operator-managed false must be preserved")
  69. }
  70. func TestEnsureSimpleModeDefaultGroups_PreservesExplicitFalse(t *testing.T) {
  71. tests := []struct {
  72. name string
  73. description string
  74. status string
  75. }{
  76. {
  77. name: "operator managed default",
  78. description: "Operator-managed group",
  79. status: service.StatusActive,
  80. },
  81. {
  82. name: "disabled auto-created default",
  83. description: simpleModeDefaultGroupDescription,
  84. status: service.StatusDisabled,
  85. },
  86. }
  87. for _, tt := range tests {
  88. t.Run(tt.name, func(t *testing.T) {
  89. ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
  90. defer cancel()
  91. client := testEntTx(t).Client()
  92. grokDefault, err := client.Group.Create().
  93. SetName(service.PlatformGrok + "-default").
  94. SetDescription(tt.description).
  95. SetPlatform(service.PlatformGrok).
  96. SetStatus(tt.status).
  97. SetSubscriptionType(service.SubscriptionTypeStandard).
  98. SetRateMultiplier(1.0).
  99. SetIsExclusive(false).
  100. SetAllowImageGeneration(false).
  101. Save(ctx)
  102. require.NoError(t, err)
  103. require.NoError(t, ensureSimpleModeDefaultGroups(ctx, client))
  104. grokDefault, err = client.Group.Get(ctx, grokDefault.ID)
  105. require.NoError(t, err)
  106. require.False(t, grokDefault.AllowImageGeneration)
  107. })
  108. }
  109. }
  110. func TestEnsureSimpleModeDefaultGroups_IgnoresSoftDeletedGroups(t *testing.T) {
  111. ctx := context.Background()
  112. tx := testEntTx(t)
  113. client := tx.Client()
  114. seedCtx, cancel := context.WithTimeout(ctx, 10*time.Second)
  115. defer cancel()
  116. // Create and then soft-delete an anthropic default group.
  117. g, err := client.Group.Create().
  118. SetName(service.PlatformAnthropic + "-default").
  119. SetPlatform(service.PlatformAnthropic).
  120. SetStatus(service.StatusActive).
  121. SetSubscriptionType(service.SubscriptionTypeStandard).
  122. SetRateMultiplier(1.0).
  123. SetIsExclusive(false).
  124. Save(seedCtx)
  125. require.NoError(t, err)
  126. _, err = client.Group.Delete().Where(group.IDEQ(g.ID)).Exec(seedCtx)
  127. require.NoError(t, err)
  128. require.NoError(t, ensureSimpleModeDefaultGroups(seedCtx, client))
  129. // New active one should exist.
  130. count, err := client.Group.Query().Where(group.NameEQ(service.PlatformAnthropic+"-default"), group.DeletedAtIsNil()).Count(seedCtx)
  131. require.NoError(t, err)
  132. require.Equal(t, 1, count)
  133. }
  134. func TestEnsureSimpleModeDefaultGroups_AntigravityNeedsTwoGroupsOnlyByCount(t *testing.T) {
  135. ctx := context.Background()
  136. tx := testEntTx(t)
  137. client := tx.Client()
  138. seedCtx, cancel := context.WithTimeout(ctx, 10*time.Second)
  139. defer cancel()
  140. mustCreateGroup(t, client, &service.Group{Name: "ag-custom-1-" + time.Now().Format(time.RFC3339Nano), Platform: service.PlatformAntigravity})
  141. mustCreateGroup(t, client, &service.Group{Name: "ag-custom-2-" + time.Now().Format(time.RFC3339Nano), Platform: service.PlatformAntigravity})
  142. require.NoError(t, ensureSimpleModeDefaultGroups(seedCtx, client))
  143. count, err := client.Group.Query().Where(group.PlatformEQ(service.PlatformAntigravity), group.DeletedAtIsNil()).Count(seedCtx)
  144. require.NoError(t, err)
  145. require.GreaterOrEqual(t, count, 2)
  146. }