client_group_reset_test.go 7.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229
  1. package service
  2. import (
  3. "testing"
  4. "github.com/mhsanaei/3x-ui/v3/internal/database"
  5. "github.com/mhsanaei/3x-ui/v3/internal/database/model"
  6. "github.com/mhsanaei/3x-ui/v3/internal/xray"
  7. )
  8. func groupByName(t *testing.T, svc *ClientService, name string) GroupSummary {
  9. t.Helper()
  10. rows, err := svc.ListGroups()
  11. if err != nil {
  12. t.Fatalf("ListGroups: %v", err)
  13. }
  14. for _, g := range rows {
  15. if g.Name == name {
  16. return g
  17. }
  18. }
  19. t.Fatalf("group %q not found in %v", name, rows)
  20. return GroupSummary{}
  21. }
  22. func seedGroupedClient(t *testing.T, email, group string, up, down int64) {
  23. t.Helper()
  24. if err := database.GetDB().Create(&model.ClientRecord{Email: email, Enable: true, Group: group}).Error; err != nil {
  25. t.Fatalf("seed client record %q: %v", email, err)
  26. }
  27. seedClientRow(t, email, 1, up, down, 0)
  28. }
  29. func TestResetGroupTraffic_ZeroesGroupButKeepsClients(t *testing.T) {
  30. initTrafficTestDB(t)
  31. svc := &ClientService{}
  32. seedGroupedClient(t, "alice", "vip", 100, 200)
  33. seedGroupedClient(t, "bob", "vip", 50, 50)
  34. before := groupByName(t, svc, "vip")
  35. if before.Up != 150 || before.Down != 250 || before.TrafficUsed != 400 || before.ClientCount != 2 {
  36. t.Fatalf("before reset: got %+v, want up=150 down=250 used=400 count=2", before)
  37. }
  38. if err := svc.ResetGroupTraffic("vip"); err != nil {
  39. t.Fatalf("ResetGroupTraffic: %v", err)
  40. }
  41. after := groupByName(t, svc, "vip")
  42. if after.Up != 0 || after.Down != 0 || after.TrafficUsed != 0 {
  43. t.Fatalf("after reset: got %+v, want up=0 down=0 used=0", after)
  44. }
  45. if after.ClientCount != 2 {
  46. t.Fatalf("after reset: client count changed to %d, want 2", after.ClientCount)
  47. }
  48. var alice xray.ClientTraffic
  49. if err := database.GetDB().Where("email = ?", "alice").First(&alice).Error; err != nil {
  50. t.Fatalf("load alice traffic: %v", err)
  51. }
  52. if alice.Up != 100 || alice.Down != 200 {
  53. t.Fatalf("client counter modified by group reset: alice up=%d down=%d, want 100/200", alice.Up, alice.Down)
  54. }
  55. }
  56. func TestResetGroupTraffic_NewTrafficAccumulatesAboveBaseline(t *testing.T) {
  57. initTrafficTestDB(t)
  58. svc := &ClientService{}
  59. seedGroupedClient(t, "carol", "team", 100, 100)
  60. if err := svc.ResetGroupTraffic("team"); err != nil {
  61. t.Fatalf("ResetGroupTraffic: %v", err)
  62. }
  63. if g := groupByName(t, svc, "team"); g.Up != 0 || g.Down != 0 {
  64. t.Fatalf("after reset: got %+v, want up=0 down=0", g)
  65. }
  66. if err := database.GetDB().Table("client_traffics").
  67. Where("email = ?", "carol").
  68. Updates(map[string]any{"up": 130, "down": 100}).Error; err != nil {
  69. t.Fatalf("bump carol traffic: %v", err)
  70. }
  71. g := groupByName(t, svc, "team")
  72. if g.Up != 30 || g.Down != 0 || g.TrafficUsed != 30 {
  73. t.Fatalf("post-bump: got %+v, want up=30 down=0 used=30", g)
  74. }
  75. }
  76. func TestResetGroupTraffic_CreatesRowForDerivedGroup(t *testing.T) {
  77. initTrafficTestDB(t)
  78. svc := &ClientService{}
  79. seedGroupedClient(t, "dave", "adhoc", 70, 30)
  80. var rows int64
  81. if err := database.GetDB().Model(&model.ClientGroup{}).Where("name = ?", "adhoc").Count(&rows).Error; err != nil {
  82. t.Fatalf("count client_groups: %v", err)
  83. }
  84. if rows != 0 {
  85. t.Fatalf("precondition: derived group should have no client_groups row, got %d", rows)
  86. }
  87. if err := svc.ResetGroupTraffic("adhoc"); err != nil {
  88. t.Fatalf("ResetGroupTraffic: %v", err)
  89. }
  90. var stored model.ClientGroup
  91. if err := database.GetDB().Where("name = ?", "adhoc").First(&stored).Error; err != nil {
  92. t.Fatalf("client_groups row not created: %v", err)
  93. }
  94. if stored.ResetUp != 70 || stored.ResetDown != 30 {
  95. t.Fatalf("baseline not snapshotted: got up=%d down=%d, want 70/30", stored.ResetUp, stored.ResetDown)
  96. }
  97. if g := groupByName(t, svc, "adhoc"); g.Up != 0 || g.Down != 0 {
  98. t.Fatalf("after reset: got %+v, want up=0 down=0", g)
  99. }
  100. }
  101. func TestResetGroupTraffic_EmptyNameRejected(t *testing.T) {
  102. initTrafficTestDB(t)
  103. svc := &ClientService{}
  104. if err := svc.ResetGroupTraffic(" "); err == nil {
  105. t.Fatal("ResetGroupTraffic(blank) = nil, want error")
  106. }
  107. }
  108. func TestGroupTotalsSurviveSingleClientReset(t *testing.T) {
  109. initTrafficTestDB(t)
  110. csvc := &ClientService{}
  111. isvc := &InboundService{}
  112. seedGroupedClient(t, "erin", "gold", 100, 200)
  113. seedGroupedClient(t, "frank", "gold", 40, 60)
  114. if err := isvc.ResetClientTrafficByEmail("erin"); err != nil {
  115. t.Fatalf("ResetClientTrafficByEmail: %v", err)
  116. }
  117. g := groupByName(t, csvc, "gold")
  118. if g.Up != 140 || g.Down != 260 || g.TrafficUsed != 400 {
  119. t.Fatalf("group totals changed by client reset: got %+v, want up=140 down=260 used=400", g)
  120. }
  121. var erin xray.ClientTraffic
  122. if err := database.GetDB().Where("email = ?", "erin").First(&erin).Error; err != nil {
  123. t.Fatalf("load erin traffic: %v", err)
  124. }
  125. if erin.Up != 0 || erin.Down != 0 {
  126. t.Fatalf("client traffic not reset: up=%d down=%d", erin.Up, erin.Down)
  127. }
  128. }
  129. func TestGroupTotalsSurviveBulkClientReset(t *testing.T) {
  130. initTrafficTestDB(t)
  131. csvc := &ClientService{}
  132. isvc := &InboundService{}
  133. seedGroupedClient(t, "gina", "silver", 10, 20)
  134. seedGroupedClient(t, "hank", "silver", 30, 40)
  135. affected, err := csvc.BulkResetTraffic(isvc, []string{"gina", "hank"})
  136. if err != nil {
  137. t.Fatalf("BulkResetTraffic: %v", err)
  138. }
  139. if affected != 2 {
  140. t.Fatalf("BulkResetTraffic affected = %d, want 2", affected)
  141. }
  142. g := groupByName(t, csvc, "silver")
  143. if g.Up != 40 || g.Down != 60 || g.TrafficUsed != 100 {
  144. t.Fatalf("group totals changed by bulk reset: got %+v, want up=40 down=60 used=100", g)
  145. }
  146. }
  147. func TestGroupTotalsSurviveClientDelete(t *testing.T) {
  148. initTrafficTestDB(t)
  149. csvc := &ClientService{}
  150. isvc := &InboundService{}
  151. seedGroupedClient(t, "iris", "bronze", 70, 30)
  152. seedGroupedClient(t, "jack", "bronze", 5, 5)
  153. var rec model.ClientRecord
  154. if err := database.GetDB().Where("email = ?", "iris").First(&rec).Error; err != nil {
  155. t.Fatalf("load iris record: %v", err)
  156. }
  157. if _, err := csvc.Delete(isvc, rec.Id, false); err != nil {
  158. t.Fatalf("Delete: %v", err)
  159. }
  160. g := groupByName(t, csvc, "bronze")
  161. if g.Up != 75 || g.Down != 35 || g.TrafficUsed != 110 {
  162. t.Fatalf("group totals changed by client delete: got %+v, want up=75 down=35 used=110", g)
  163. }
  164. if g.ClientCount != 1 {
  165. t.Fatalf("client count = %d, want 1", g.ClientCount)
  166. }
  167. var trafficRows int64
  168. if err := database.GetDB().Model(&xray.ClientTraffic{}).Where("email = ?", "iris").Count(&trafficRows).Error; err != nil {
  169. t.Fatalf("count iris traffic rows: %v", err)
  170. }
  171. if trafficRows != 0 {
  172. t.Fatalf("iris traffic row survived delete")
  173. }
  174. }
  175. func TestGroupResetStillZeroesAfterBaselineAdjustments(t *testing.T) {
  176. initTrafficTestDB(t)
  177. csvc := &ClientService{}
  178. isvc := &InboundService{}
  179. seedGroupedClient(t, "kate", "iron", 100, 100)
  180. if err := isvc.ResetClientTrafficByEmail("kate"); err != nil {
  181. t.Fatalf("ResetClientTrafficByEmail: %v", err)
  182. }
  183. if g := groupByName(t, csvc, "iron"); g.TrafficUsed != 200 {
  184. t.Fatalf("pre group-reset: got %+v, want used=200", g)
  185. }
  186. if err := csvc.ResetGroupTraffic("iron"); err != nil {
  187. t.Fatalf("ResetGroupTraffic: %v", err)
  188. }
  189. if g := groupByName(t, csvc, "iron"); g.Up != 0 || g.Down != 0 || g.TrafficUsed != 0 {
  190. t.Fatalf("group reset did not zero adjusted baselines: got %+v", g)
  191. }
  192. }