golimiter/bucket_manager_test.go

package golimiter

import (
	"testing"
	"time"
)

func TestTokenBucketManager(t *testing.T) {
	tbManager := NewTokenBucketManager()

	tb, err := tbManager.CreateTokenBucket("127.0.0.1", 10, 5, time.Second)
	if err != nil {
		t.Fatalf("Failed to create token bucket: %v", err)
	}

	foundBucket := tbManager.GetTokenBucket("127.0.0.1")
	if foundBucket != tb {
		t.Error("Retrieved token bucket does not match created bucket")
	}

	nonExistentBucket := tbManager.GetTokenBucket("192.168.1.1")
	if nonExistentBucket != nil {
		t.Error("Expected nil for non-existent token bucket")
	}
}


func TestTokenBucketCleanup(t *testing.T) {
	tbManager := NewTokenBucketManager()
	tbManager.CreateTokenBucket("127.0.0.1", 10, 5, 2*time.Second)

	time.Sleep(3 * time.Second)

	go tbManager.StartCleanupTask(time.Second)

	// Wait for cleanup to run
	time.Sleep(2 * time.Second)

	tb := tbManager.GetTokenBucket("127.0.0.1")
	if tb != nil {
		t.Error("Expired token bucket was not cleaned up")
	}
}