github.com/observiq/carbon@v0.9.11-0.20200820160507-1b872e368a5e/operator/builtin/transformer/rate_limit_test.go (about)

     1  package transformer
     2  
     3  import (
     4  	"context"
     5  	"sync"
     6  	"testing"
     7  	"time"
     8  
     9  	"github.com/observiq/carbon/entry"
    10  	"github.com/observiq/carbon/operator"
    11  	"github.com/observiq/carbon/testutil"
    12  	"github.com/stretchr/testify/mock"
    13  	"github.com/stretchr/testify/require"
    14  )
    15  
    16  func TestRateLimit(t *testing.T) {
    17  	t.Parallel()
    18  
    19  	cfg := NewRateLimitConfig("my_rate_limit")
    20  	cfg.OutputIDs = []string{"output1"}
    21  	cfg.Burst = 1
    22  	cfg.Rate = 10
    23  
    24  	rateLimit, err := cfg.Build(testutil.NewBuildContext(t))
    25  	require.NoError(t, err)
    26  
    27  	receivedLog := make(chan struct{}, 100)
    28  	mockOutput := testutil.NewMockOperator("output1")
    29  	mockOutput.On("Process", mock.Anything, mock.Anything).Return(nil).Run(func(args mock.Arguments) {
    30  		receivedLog <- struct{}{}
    31  	})
    32  
    33  	err = rateLimit.SetOutputs([]operator.Operator{mockOutput})
    34  	require.NoError(t, err)
    35  
    36  	err = rateLimit.Start()
    37  	require.NoError(t, err)
    38  
    39  	ctx, cancel := context.WithCancel(context.Background())
    40  	wg := &sync.WaitGroup{}
    41  	wg.Add(1)
    42  	go func() {
    43  		defer wg.Done()
    44  		for {
    45  			select {
    46  			case <-ctx.Done():
    47  				return
    48  			default:
    49  				err := rateLimit.Process(ctx, entry.New())
    50  				require.NoError(t, err)
    51  			}
    52  		}
    53  	}()
    54  
    55  	i := 0
    56  	timeout := time.After(time.Second)
    57  LOOP:
    58  	for {
    59  		select {
    60  		case <-receivedLog:
    61  			i++
    62  		case <-timeout:
    63  			break LOOP
    64  		}
    65  	}
    66  
    67  	cancel()
    68  	wg.Wait()
    69  
    70  	require.InDelta(t, 10, i, 3)
    71  }