token_bucket_test.go

package equalizer

import (
	"context"
	"testing"
	"time"

	"github.com/reugn/equalizer/internal/assert"
)

func TestTokenBucket(t *testing.T) {
	tokenBucket, err := NewTokenBucket(32, 100*time.Millisecond)
	assert.IsNil(t, err)
	var quota bool
	for i := 0; i < 32; i++ {
		quota = tokenBucket.TryAcquire()
	}
	assert.Equal(t, quota, true)

	quota = tokenBucket.TryAcquire()
	assert.Equal(t, quota, false)

	time.Sleep(110 * time.Millisecond)

	quota = tokenBucket.TryAcquire()
	assert.Equal(t, quota, true)

	assert.IsNil(t, tokenBucket.Acquire(context.Background()))
}

func TestTokenBucket_Context(t *testing.T) {
	tokenBucket, err := NewTokenBucket(1, time.Second)
	assert.IsNil(t, err)
	assert.IsNil(t, tokenBucket.Acquire(context.Background()))

	ctx, cancel := context.WithDeadline(context.Background(), time.Now().Add(10*time.Millisecond))
	defer cancel()
	err = tokenBucket.Acquire(ctx)
	assert.ErrorIs(t, err, context.DeadlineExceeded)
}

func TestTokenBucket_Errors(t *testing.T) {
	tokenBucket, err := NewTokenBucket(-1, time.Second)
	assert.ErrorContains(t, err, "nonpositive capacity")
	assert.IsNil(t, tokenBucket)
}