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