client_group_reset_test.go 3.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127
  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. }