github.com/grafana/pyroscope@v1.18.0/pkg/util/ratelimit/ratelimit_test.go (about)

     1  package ratelimit
     2  
     3  import (
     4  	"testing"
     5  	"time"
     6  
     7  	"github.com/stretchr/testify/assert"
     8  	"github.com/stretchr/testify/require"
     9  )
    10  
    11  func TestLimiter(t *testing.T) {
    12  	var sleptFor []time.Duration
    13  	now := time.Unix(0, 0)
    14  
    15  	l := NewLimiter(100)
    16  	l.now = func() time.Time { return now }
    17  	l.sleep = func(d time.Duration) {
    18  		sleptFor = append(sleptFor, d)
    19  		now = now.Add(d)
    20  	}
    21  
    22  	l.Wait(100)
    23  	assert.Len(t, sleptFor, 0)
    24  
    25  	l.Wait(100)
    26  	require.Len(t, sleptFor, 1)
    27  	assert.Equal(t, time.Second, sleptFor[0])
    28  
    29  	l.Wait(100)
    30  	require.Len(t, sleptFor, 2)
    31  	assert.Equal(t, time.Second, sleptFor[1])
    32  }