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

     1  package ratelimit
     2  
     3  import (
     4  	"io"
     5  	"math/rand"
     6  	"testing"
     7  	"time"
     8  
     9  	"github.com/stretchr/testify/assert"
    10  	"github.com/stretchr/testify/require"
    11  )
    12  
    13  func TestWriter(t *testing.T) {
    14  	var sleptFor time.Duration
    15  	now := time.Unix(0, 0)
    16  
    17  	limiter := NewLimiter(100)
    18  	limiter.now = func() time.Time { return now }
    19  	limiter.sleep = func(d time.Duration) {
    20  		sleptFor += d
    21  		now = now.Add(d)
    22  	}
    23  
    24  	w := NewWriter(io.Discard, limiter)
    25  
    26  	const N = 1 << 10
    27  	n, err := io.CopyN(w, rand.New(rand.NewSource(42)), N)
    28  	require.NoError(t, err)
    29  	assert.EqualValues(t, n, N)
    30  
    31  	t.Log("written in", sleptFor)
    32  	assert.Greater(t, sleptFor, time.Second*9)
    33  }