github.com/rudderlabs/rudder-go-kit@v0.30.0/sync/group_test.go (about)

     1  package sync
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"sync/atomic"
     7  	"testing"
     8  	"time"
     9  
    10  	"github.com/stretchr/testify/require"
    11  )
    12  
    13  func TestEagerGroupWithLimit(t *testing.T) {
    14  	g, ctx := NewEagerGroup(context.Background(), 2)
    15  	var count atomic.Int64
    16  	// One of the following three goroutines should DEFINITELY NOT be executed due to the limit of 2 and the context being cancelled.
    17  	// The context should get cancelled automatically because the first two routines returned an error.
    18  	g.Go(func() error {
    19  		t.Log("one")
    20  		count.Add(1)
    21  		return fmt.Errorf("one")
    22  	})
    23  	g.Go(func() error {
    24  		t.Log("two")
    25  		count.Add(1)
    26  		return fmt.Errorf("two")
    27  	})
    28  	g.Go(func() error {
    29  		t.Log("three")
    30  		count.Add(1)
    31  		return fmt.Errorf("three")
    32  	})
    33  	require.Error(t, g.Wait(), "We expect group.Wait() to return an error")
    34  	ok := true
    35  	select {
    36  	case <-ctx.Done():
    37  		_, ok = <-ctx.Done()
    38  	case <-time.After(time.Second):
    39  	}
    40  	require.False(t, ok, "We expect the context to be cancelled")
    41  	require.True(t, 1 <= count.Load() && count.Load() <= 2, "We expect count to be between 1 and 2")
    42  }
    43  
    44  func TestEagerGroupWithNoLimit(t *testing.T) {
    45  	ctx, cancel := context.WithCancel(context.Background())
    46  	g, ctx := NewEagerGroup(ctx, 0)
    47  	funcCounter := &atomic.Int64{}
    48  
    49  	go func() {
    50  		for {
    51  			if funcCounter.Load() > 10 {
    52  				cancel()
    53  				return
    54  			}
    55  		}
    56  	}()
    57  
    58  	for i := 0; i < 10000; i++ {
    59  		g.Go(func() error {
    60  			select {
    61  			case <-ctx.Done():
    62  				return ctx.Err()
    63  			default:
    64  			}
    65  			funcCounter.Add(1)
    66  			return nil
    67  		})
    68  	}
    69  	require.ErrorIs(t, g.Wait(), ctx.Err(), "We expect group.Wait() to return the context error")
    70  	_, ok := <-ctx.Done()
    71  	require.False(t, ok, "We expect the context to be cancelled")
    72  	t.Log(funcCounter.Load(), "funcs executed")
    73  	// We expect between 10 and 10000 funcs to be executed
    74  	// because group tries to return early if context is cancelled
    75  	require.Less(
    76  		t,
    77  		funcCounter.Load(),
    78  		int64(10000),
    79  		"Expected less than 1000 funcs to be executed",
    80  	)
    81  }
    82  
    83  func TestNoInitEagerGroup(t *testing.T) {
    84  	g := &EagerGroup{}
    85  	f := func() error { return nil }
    86  	require.Panics(
    87  		t,
    88  		func() { g.Go(f) },
    89  		"We expect a panic when calling Go on a group that has not been initialized with NewEagerGroup",
    90  	)
    91  }