github.com/Jeffail/benthos/v3@v3.65.0/internal/component/ratelimit/rate_limit_v2_test.go (about)

     1  package ratelimit
     2  
     3  import (
     4  	"context"
     5  	"testing"
     6  	"time"
     7  
     8  	"github.com/Jeffail/benthos/v3/lib/metrics"
     9  	"github.com/stretchr/testify/assert"
    10  )
    11  
    12  type closableRateLimit struct {
    13  	closed bool
    14  }
    15  
    16  func (c *closableRateLimit) Access(ctx context.Context) (time.Duration, error) {
    17  	return 0, nil
    18  }
    19  
    20  func (c *closableRateLimit) Close(ctx context.Context) error {
    21  	c.closed = true
    22  	return nil
    23  }
    24  
    25  func TestRateLimitAirGapShutdown(t *testing.T) {
    26  	rl := &closableRateLimit{}
    27  	agrl := NewV2ToV1RateLimit(rl, metrics.Noop())
    28  
    29  	err := agrl.WaitForClose(time.Millisecond * 5)
    30  	assert.EqualError(t, err, "action timed out")
    31  	assert.False(t, rl.closed)
    32  
    33  	agrl.CloseAsync()
    34  	err = agrl.WaitForClose(time.Millisecond * 5)
    35  	assert.NoError(t, err)
    36  	assert.True(t, rl.closed)
    37  }