github.com/onflow/flow-go@v0.35.7-crescendo-preview.23-atree-inlining/engine/broadcaster_test.go (about)

     1  package engine_test
     2  
     3  import (
     4  	"context"
     5  	"sync"
     6  	"testing"
     7  	"time"
     8  
     9  	"github.com/stretchr/testify/assert"
    10  	"go.uber.org/atomic"
    11  
    12  	"github.com/onflow/flow-go/engine"
    13  	"github.com/onflow/flow-go/utils/unittest"
    14  )
    15  
    16  func TestPublish(t *testing.T) {
    17  	t.Parallel()
    18  
    19  	t.Run("no subscribers", func(t *testing.T) {
    20  		t.Parallel()
    21  		b := engine.NewBroadcaster()
    22  		unittest.RequireReturnsBefore(t, b.Publish, 100*time.Millisecond, "publish never finished")
    23  	})
    24  
    25  	t.Run("all subscribers notified", func(t *testing.T) {
    26  		t.Parallel()
    27  		notifierCount := 10
    28  		recievedCount := atomic.NewInt32(0)
    29  
    30  		b := engine.NewBroadcaster()
    31  
    32  		// setup subscribers to listen for a notification then return
    33  		subscribers := sync.WaitGroup{}
    34  		subscribers.Add(notifierCount)
    35  
    36  		for i := 0; i < notifierCount; i++ {
    37  			notifier := engine.NewNotifier()
    38  			b.Subscribe(notifier)
    39  			go func() {
    40  				defer subscribers.Done()
    41  				<-notifier.Channel()
    42  				recievedCount.Inc()
    43  			}()
    44  		}
    45  
    46  		b.Publish()
    47  
    48  		unittest.RequireReturnsBefore(t, subscribers.Wait, 100*time.Millisecond, "wait never finished")
    49  
    50  		// there should be one notification for each subscriber
    51  		assert.Equal(t, int32(notifierCount), recievedCount.Load())
    52  	})
    53  
    54  	t.Run("all subscribers notified at least once", func(t *testing.T) {
    55  		t.Parallel()
    56  		notifierCount := 10
    57  		notifiedCounts := make([]int, notifierCount)
    58  
    59  		ctx, cancel := context.WithCancel(context.Background())
    60  
    61  		b := engine.NewBroadcaster()
    62  
    63  		// setup subscribers to listen for notifications until the context is cancelled
    64  		subscribers := sync.WaitGroup{}
    65  		subscribers.Add(notifierCount)
    66  
    67  		for i := 0; i < notifierCount; i++ {
    68  			notifier := engine.NewNotifier()
    69  			b.Subscribe(notifier)
    70  
    71  			go func(i int) {
    72  				defer subscribers.Done()
    73  
    74  				for {
    75  					select {
    76  					case <-ctx.Done():
    77  						return
    78  					case <-notifier.Channel():
    79  						notifiedCounts[i]++
    80  					}
    81  				}
    82  			}(i)
    83  		}
    84  
    85  		// setup publisher to publish notifications concurrently
    86  		publishers := sync.WaitGroup{}
    87  		publishers.Add(20)
    88  
    89  		for i := 0; i < 20; i++ {
    90  			go func() {
    91  				defer publishers.Done()
    92  				b.Publish()
    93  
    94  				// pause to allow the scheduler to switch to another goroutine
    95  				time.Sleep(time.Millisecond)
    96  			}()
    97  		}
    98  
    99  		// wait for publishers to finish, then cancel subscribers' context
   100  		unittest.RequireReturnsBefore(t, publishers.Wait, 100*time.Millisecond, "publishers never finished")
   101  		time.Sleep(100 * time.Millisecond)
   102  
   103  		cancel()
   104  
   105  		unittest.RequireReturnsBefore(t, subscribers.Wait, 100*time.Millisecond, "receivers never finished")
   106  
   107  		// all subscribers should have been notified at least once
   108  		for i, count := range notifiedCounts {
   109  			assert.GreaterOrEqualf(t, count, 1, "notifier %d was not notified", i)
   110  		}
   111  	})
   112  }