github.com/olljanat/moby@v1.13.1/pkg/pubsub/publisher_test.go (about)

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