publisher_test.go 2.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142
  1. package pubsub
  2. import (
  3. "fmt"
  4. "testing"
  5. "time"
  6. )
  7. func TestSendToOneSub(t *testing.T) {
  8. p := NewPublisher(100*time.Millisecond, 10)
  9. c := p.Subscribe()
  10. p.Publish("hi")
  11. msg := <-c
  12. if msg.(string) != "hi" {
  13. t.Fatalf("expected message hi but received %v", msg)
  14. }
  15. }
  16. func TestSendToMultipleSubs(t *testing.T) {
  17. p := NewPublisher(100*time.Millisecond, 10)
  18. subs := []chan interface{}{}
  19. subs = append(subs, p.Subscribe(), p.Subscribe(), p.Subscribe())
  20. p.Publish("hi")
  21. for _, c := range subs {
  22. msg := <-c
  23. if msg.(string) != "hi" {
  24. t.Fatalf("expected message hi but received %v", msg)
  25. }
  26. }
  27. }
  28. func TestEvictOneSub(t *testing.T) {
  29. p := NewPublisher(100*time.Millisecond, 10)
  30. s1 := p.Subscribe()
  31. s2 := p.Subscribe()
  32. p.Evict(s1)
  33. p.Publish("hi")
  34. if _, ok := <-s1; ok {
  35. t.Fatal("expected s1 to not receive the published message")
  36. }
  37. msg := <-s2
  38. if msg.(string) != "hi" {
  39. t.Fatalf("expected message hi but received %v", msg)
  40. }
  41. }
  42. func TestClosePublisher(t *testing.T) {
  43. p := NewPublisher(100*time.Millisecond, 10)
  44. subs := []chan interface{}{}
  45. subs = append(subs, p.Subscribe(), p.Subscribe(), p.Subscribe())
  46. p.Close()
  47. for _, c := range subs {
  48. if _, ok := <-c; ok {
  49. t.Fatal("expected all subscriber channels to be closed")
  50. }
  51. }
  52. }
  53. const sampleText = "test"
  54. type testSubscriber struct {
  55. dataCh chan interface{}
  56. ch chan error
  57. }
  58. func (s *testSubscriber) Wait() error {
  59. return <-s.ch
  60. }
  61. func newTestSubscriber(p *Publisher) *testSubscriber {
  62. ts := &testSubscriber{
  63. dataCh: p.Subscribe(),
  64. ch: make(chan error),
  65. }
  66. go func() {
  67. for data := range ts.dataCh {
  68. s, ok := data.(string)
  69. if !ok {
  70. ts.ch <- fmt.Errorf("Unexpected type %T", data)
  71. break
  72. }
  73. if s != sampleText {
  74. ts.ch <- fmt.Errorf("Unexpected text %s", s)
  75. break
  76. }
  77. }
  78. close(ts.ch)
  79. }()
  80. return ts
  81. }
  82. // for testing with -race
  83. func TestPubSubRace(t *testing.T) {
  84. p := NewPublisher(0, 1024)
  85. var subs [](*testSubscriber)
  86. for j := 0; j < 50; j++ {
  87. subs = append(subs, newTestSubscriber(p))
  88. }
  89. for j := 0; j < 1000; j++ {
  90. p.Publish(sampleText)
  91. }
  92. time.AfterFunc(1*time.Second, func() {
  93. for _, s := range subs {
  94. p.Evict(s.dataCh)
  95. }
  96. })
  97. for _, s := range subs {
  98. s.Wait()
  99. }
  100. }
  101. func BenchmarkPubSub(b *testing.B) {
  102. for i := 0; i < b.N; i++ {
  103. b.StopTimer()
  104. p := NewPublisher(0, 1024)
  105. var subs [](*testSubscriber)
  106. for j := 0; j < 50; j++ {
  107. subs = append(subs, newTestSubscriber(p))
  108. }
  109. b.StartTimer()
  110. for j := 0; j < 1000; j++ {
  111. p.Publish(sampleText)
  112. }
  113. time.AfterFunc(1*time.Second, func() {
  114. for _, s := range subs {
  115. p.Evict(s.dataCh)
  116. }
  117. })
  118. for _, s := range subs {
  119. if err := s.Wait(); err != nil {
  120. b.Fatal(err)
  121. }
  122. }
  123. }
  124. }