golimiter/golimiter_test.go

package golimiter

import (
	"testing"
	"time"
)

func TestTokenBucket(t *testing.T) {
	tb := NewTokenBucket(10, 5, time.Second)

	for i := 0; i < 10; i++ {
		if !tb.TakeToken() {
			t.Errorf("Expected token %d, but bucket is empty", i)
		}
	}

	if tb.TakeToken() {
		t.Errorf("Expected bucket to be empty, but token was taken")
	}
}