Improve test coverage

bucket/leaky.go

@@ -13,20 +13,28 @@ type LeakyLimiter struct {
 
 	capacity, level, rate float64
 	lastUpdatedAt         time.Time
+	clock                 clock
 }
 
 // NewLeakyLimiter creates a new leaky bucket limiter.
 // Capacity is the maximum bucket size. Rate is how many requests drain per second.
 func NewLeakyLimiter(capacity, rate uint32) *LeakyLimiter {
+	return NewLeakyLimiterWithClock(capacity, rate, realClock{})
+}
+
+// NewLeakyLimiterWithClock creates a new leaky bucket limiter with a custom clock.
+// Use this constructor for testing with a mock clock.
+func NewLeakyLimiterWithClock(capacity, rate uint32, clock clock) *LeakyLimiter {
 	return &LeakyLimiter{
 		capacity:      float64(capacity),
 		rate:          float64(rate),
-		lastUpdatedAt: time.Now(),
+		clock:         clock,
+		lastUpdatedAt: clock.Now(),
 	}
 }
 
 func (lim *LeakyLimiter) update() {
-	t := time.Now()
+	t := lim.clock.Now()
 	if t.Before(lim.lastUpdatedAt) {
 		return
 	}

bucket/leaky_test.go

@@ -4,9 +4,11 @@ import (
 	"sync"
 	"sync/atomic"
 	"testing"
+	"time"
 
 	"github.com/serroba/rate/bucket"
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 )
 
 func TestLeakyLimiter_Allow(t *testing.T) {
@@ -70,3 +72,57 @@ func TestLeakyLimiter_Allow_Concurrent(t *testing.T) {
 	assert.Equal(t, int64(10), allow.Load())
 	assert.Equal(t, int64(5), deny.Load())
 }
+
+func TestLeakyLimiter_Allow_ClockGoesBackwards(t *testing.T) {
+	t.Parallel()
+
+	clock := &testClock{now: time.Now()}
+	lim := bucket.NewLeakyLimiterWithClock(1, 1, clock)
+
+	// Fill the bucket
+	require.True(t, lim.Allow())
+	require.False(t, lim.Allow())
+
+	// Move clock backwards - should not drain
+	clock.now = clock.now.Add(-1 * time.Second)
+
+	require.False(t, lim.Allow())
+}
+
+func TestLeakyLimiter_Allow_Drains(t *testing.T) {
+	t.Parallel()
+
+	clock := &testClock{now: time.Now()}
+	lim := bucket.NewLeakyLimiterWithClock(2, 2, clock) // drains 2 per second
+
+	// Fill the bucket
+	require.True(t, lim.Allow())
+	require.True(t, lim.Allow())
+	require.False(t, lim.Allow())
+
+	// Advance 1 second - should drain 2, bucket now empty
+	clock.advance(1 * time.Second)
+
+	// Can fill again
+	require.True(t, lim.Allow())
+	require.True(t, lim.Allow())
+	require.False(t, lim.Allow())
+}
+
+func TestLeakyLimiter_Allow_PartialDrain(t *testing.T) {
+	t.Parallel()
+
+	clock := &testClock{now: time.Now()}
+	lim := bucket.NewLeakyLimiterWithClock(2, 2, clock) // drains 2 per second
+
+	// Fill the bucket
+	require.True(t, lim.Allow())
+	require.True(t, lim.Allow())
+	require.False(t, lim.Allow())
+
+	// Advance 0.5 seconds - should drain 1, leaving room for 1
+	clock.advance(500 * time.Millisecond)
+
+	require.True(t, lim.Allow())
+	require.False(t, lim.Allow())
+}

registry/registry_test.go

@@ -21,48 +21,48 @@ type StrategyConfig interface {
 
 // TokenBucketConfig configures a token bucket rate limiter.
 type TokenBucketConfig struct {
-	Capacity uint32
-	Rate     uint32
+	capacity uint32
+	rate     uint32
 }
 
 func (c TokenBucketConfig) Name() string { return "token bucket" }
 
 func (c TokenBucketConfig) Build() registry.LimiterFactory {
 	return func() registry.Limiter {
-		return bucket.NewTokenLimiter(c.Capacity, c.Rate)
+		return bucket.NewTokenLimiter(c.capacity, c.rate)
 	}
 }
 
 // LeakyBucketConfig configures a leaky bucket rate limiter.
 type LeakyBucketConfig struct {
-	Capacity uint32
-	Rate     uint32
+	capacity uint32
+	rate     uint32
 }
 
 func (c LeakyBucketConfig) Name() string { return "leaky bucket" }
 
 func (c LeakyBucketConfig) Build() registry.LimiterFactory {
 	return func() registry.Limiter {
-		return bucket.NewLeakyLimiter(c.Capacity, c.Rate)
+		return bucket.NewLeakyLimiter(c.capacity, c.rate)
 	}
 }
 
 // FixedWindowConfig configures a fixed window rate limiter.
 type FixedWindowConfig struct {
-	Limit  uint32
-	Window time.Duration
+	limit  uint32
+	window time.Duration
 }
 
 func (c FixedWindowConfig) Name() string { return "fixed window" }
 
 func (c FixedWindowConfig) Build() registry.LimiterFactory {
-	w := c.Window
+	w := c.window
 	if w == 0 {
 		w = time.Hour
 	}
 
 	return func() registry.Limiter {
-		return window.NewFixedLimiter(c.Limit, w)
+		return window.NewFixedLimiter(c.limit, w)
 	}
 }
 
@@ -82,9 +82,9 @@ func (s SlidingWindowConfig) Build() registry.LimiterFactory {
 // Helper to create all strategies with given parameters.
 func allStrategies(capacity uint32, rate uint32, win time.Duration) []StrategyConfig {
 	return []StrategyConfig{
-		TokenBucketConfig{Capacity: capacity, Rate: rate},
-		LeakyBucketConfig{Capacity: capacity, Rate: rate},
-		FixedWindowConfig{Limit: capacity, Window: win},
+		TokenBucketConfig{capacity: capacity, rate: rate},
+		LeakyBucketConfig{capacity: capacity, rate: rate},
+		FixedWindowConfig{limit: capacity, window: win},
 		SlidingWindowConfig{limit: capacity, duration: win},
 	}
 }