github.com/iqoqo/nomad@v0.11.3-0.20200911112621-d7021c74d101/drivers/shared/eventer/eventer_test.go (about)

     1  package eventer
     2  
     3  import (
     4  	"context"
     5  	"sync"
     6  	"testing"
     7  	"time"
     8  
     9  	"github.com/hashicorp/nomad/helper/testlog"
    10  	"github.com/hashicorp/nomad/plugins/drivers"
    11  	"github.com/stretchr/testify/require"
    12  )
    13  
    14  func TestEventer(t *testing.T) {
    15  	t.Parallel()
    16  	require := require.New(t)
    17  
    18  	ctx, cancel := context.WithCancel(context.Background())
    19  	defer cancel()
    20  	e := NewEventer(ctx, testlog.HCLogger(t))
    21  
    22  	events := []*drivers.TaskEvent{
    23  		{
    24  			TaskID:    "a",
    25  			Timestamp: time.Now(),
    26  		},
    27  		{
    28  			TaskID:    "b",
    29  			Timestamp: time.Now(),
    30  		},
    31  		{
    32  			TaskID:    "c",
    33  			Timestamp: time.Now(),
    34  		},
    35  	}
    36  
    37  	ctx1, cancel1 := context.WithCancel(context.Background())
    38  	defer cancel1()
    39  	consumer1, err := e.TaskEvents(ctx1)
    40  	require.NoError(err)
    41  	ctx2 := (context.Background())
    42  	consumer2, err := e.TaskEvents(ctx2)
    43  	require.NoError(err)
    44  
    45  	var buffer1, buffer2 []*drivers.TaskEvent
    46  	var wg sync.WaitGroup
    47  	wg.Add(2)
    48  	go func() {
    49  		defer wg.Done()
    50  		var i int
    51  		for event := range consumer1 {
    52  			i++
    53  			buffer1 = append(buffer1, event)
    54  			if i == len(events) {
    55  				return
    56  			}
    57  		}
    58  	}()
    59  	go func() {
    60  		defer wg.Done()
    61  		var i int
    62  		for event := range consumer2 {
    63  			i++
    64  			buffer2 = append(buffer2, event)
    65  			if i == len(events) {
    66  				return
    67  			}
    68  		}
    69  	}()
    70  
    71  	for _, event := range events {
    72  		require.NoError(e.EmitEvent(event))
    73  	}
    74  
    75  	wg.Wait()
    76  	require.Exactly(events, buffer1)
    77  	require.Exactly(events, buffer2)
    78  }
    79  
    80  func TestEventer_iterateConsumers(t *testing.T) {
    81  	t.Parallel()
    82  	require := require.New(t)
    83  
    84  	e := &Eventer{
    85  		events: make(chan *drivers.TaskEvent),
    86  		ctx:    context.Background(),
    87  		logger: testlog.HCLogger(t),
    88  	}
    89  
    90  	ev := &drivers.TaskEvent{
    91  		TaskID:    "a",
    92  		Timestamp: time.Now(),
    93  	}
    94  
    95  	ctx1, cancel1 := context.WithCancel(context.Background())
    96  	consumer, err := e.TaskEvents(ctx1)
    97  	require.NoError(err)
    98  	require.Equal(1, len(e.consumers))
    99  
   100  	var wg sync.WaitGroup
   101  	wg.Add(1)
   102  	go func() {
   103  		defer wg.Done()
   104  		ev1, ok := <-consumer
   105  		require.Exactly(ev, ev1)
   106  		require.True(ok)
   107  	}()
   108  	e.iterateConsumers(ev)
   109  	wg.Wait()
   110  
   111  	go func() {
   112  		cancel1()
   113  		e.iterateConsumers(ev)
   114  	}()
   115  	ev1, ok := <-consumer
   116  	require.False(ok)
   117  	require.Nil(ev1)
   118  	require.Equal(0, len(e.consumers))
   119  }