github.com/Jeffail/benthos/v3@v3.65.0/public/service/rate_limit_test.go (about)

     1  package service
     2  
     3  import (
     4  	"context"
     5  	"errors"
     6  	"testing"
     7  	"time"
     8  
     9  	"github.com/Jeffail/benthos/v3/lib/metrics"
    10  	"github.com/stretchr/testify/assert"
    11  )
    12  
    13  type closableRateLimit struct {
    14  	next   time.Duration
    15  	err    error
    16  	closed bool
    17  }
    18  
    19  func (c *closableRateLimit) Access(ctx context.Context) (time.Duration, error) {
    20  	return c.next, c.err
    21  }
    22  
    23  func (c *closableRateLimit) Close(ctx context.Context) error {
    24  	c.closed = true
    25  	return nil
    26  }
    27  
    28  func TestRateLimitAirGapShutdown(t *testing.T) {
    29  	rl := &closableRateLimit{
    30  		next: time.Second,
    31  	}
    32  	agrl := newAirGapRateLimit(rl, metrics.Noop())
    33  
    34  	tout, err := agrl.Access()
    35  	assert.NoError(t, err)
    36  	assert.Equal(t, time.Second, tout)
    37  
    38  	rl.next = time.Millisecond
    39  	rl.err = errors.New("test error")
    40  
    41  	tout, err = agrl.Access()
    42  	assert.EqualError(t, err, "test error")
    43  	assert.Equal(t, time.Millisecond, tout)
    44  
    45  	err = agrl.WaitForClose(time.Millisecond * 5)
    46  	assert.EqualError(t, err, "action timed out")
    47  	assert.False(t, rl.closed)
    48  
    49  	agrl.CloseAsync()
    50  	err = agrl.WaitForClose(time.Millisecond * 5)
    51  	assert.NoError(t, err)
    52  	assert.True(t, rl.closed)
    53  }
    54  
    55  //------------------------------------------------------------------------------
    56  
    57  type closableRateLimitType struct {
    58  	next   time.Duration
    59  	err    error
    60  	closed bool
    61  }
    62  
    63  func (c *closableRateLimitType) Access() (time.Duration, error) {
    64  	return c.next, c.err
    65  }
    66  
    67  func (c *closableRateLimitType) CloseAsync() {
    68  	c.closed = true
    69  }
    70  
    71  func (c *closableRateLimitType) WaitForClose(tout time.Duration) error {
    72  	return nil
    73  }
    74  
    75  func TestRateLimitReverseAirGapShutdown(t *testing.T) {
    76  	rl := &closableRateLimitType{
    77  		next: time.Second,
    78  	}
    79  	agrl := newReverseAirGapRateLimit(rl)
    80  
    81  	tout, err := agrl.Access(context.Background())
    82  	assert.NoError(t, err)
    83  	assert.Equal(t, time.Second, tout)
    84  
    85  	rl.next = time.Millisecond
    86  	rl.err = errors.New("test error")
    87  
    88  	tout, err = agrl.Access(context.Background())
    89  	assert.EqualError(t, err, "test error")
    90  	assert.Equal(t, time.Millisecond, tout)
    91  
    92  	assert.NoError(t, agrl.Close(context.Background()))
    93  	assert.True(t, rl.closed)
    94  }