reset_writer_test.go 1.7 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768
  1. package service
  2. import (
  3. "context"
  4. "testing"
  5. "time"
  6. "github.com/mhsanaei/3x-ui/v3/internal/database/model"
  7. "github.com/mhsanaei/3x-ui/v3/internal/web/runtime"
  8. )
  9. type blockingResetRuntime struct {
  10. fakeNodeRuntime
  11. reached chan struct{}
  12. release chan struct{}
  13. }
  14. func (b *blockingResetRuntime) ResetAllTraffics(context.Context) error {
  15. close(b.reached)
  16. <-b.release
  17. return nil
  18. }
  19. func TestResetAllTrafficsDoesNotBlockWriterOnNodeCall(t *testing.T) {
  20. db := initTrafficTestDB(t)
  21. resetTrafficWriterForTest(t)
  22. StartTrafficWriter()
  23. mgr := runtime.NewManager(runtime.LocalDeps{APIPort: func() int { return 0 }})
  24. runtime.SetManager(mgr)
  25. t.Cleanup(func() { runtime.SetManager(nil) })
  26. node := &model.Node{Name: "n1", Address: "127.0.0.1", Port: 2096, ApiToken: "tok", Enable: true, Status: "online"}
  27. if err := db.Create(node).Error; err != nil {
  28. t.Fatalf("create node: %v", err)
  29. }
  30. fake := &blockingResetRuntime{reached: make(chan struct{}), release: make(chan struct{})}
  31. mgr.SetRuntimeOverride(node.Id, fake)
  32. done := make(chan error, 1)
  33. go func() { done <- (&InboundService{}).ResetAllTraffics() }()
  34. select {
  35. case <-fake.reached:
  36. case <-time.After(3 * time.Second):
  37. close(fake.release)
  38. t.Fatal("node ResetAllTraffics was never reached")
  39. }
  40. writerFree := make(chan error, 1)
  41. go func() { writerFree <- submitTrafficWrite(func() error { return nil }) }()
  42. select {
  43. case err := <-writerFree:
  44. if err != nil {
  45. close(fake.release)
  46. t.Fatalf("concurrent writer submit failed: %v", err)
  47. }
  48. case <-time.After(2 * time.Second):
  49. close(fake.release)
  50. <-done
  51. t.Fatal("the serial traffic writer was blocked by a node reset HTTP call")
  52. }
  53. close(fake.release)
  54. if err := <-done; err != nil {
  55. t.Fatalf("ResetAllTraffics: %v", err)
  56. }
  57. }