package middleware

import (
	"testing"
	"time"

	"golang.org/x/time/rate"
)

// TestAddressLimiterDropsOnlyFullBuckets checks the sweep in allow:
// a minute after the last one, it drops an address whose bucket has
// filled up again, and keeps one still short of tokens, whose limit
// would otherwise start over.
func TestAddressLimiterDropsOnlyFullBuckets(t *testing.T) {
	t.Parallel()

	const (
		refilled = "198.51.100.1"
		drained  = "198.51.100.2"
	)

	// Two a minute: one token back every 30 seconds.
	limiter := &addressLimiter{
		burst:  2,
		byAddr: make(map[string]*rate.Limiter),
		limit:  rate.Every(30 * time.Second),
	}

	start := time.Now()

	// The first call sweeps the empty map and takes one of two tokens.
	limiter.allow(refilled, start)
	limiter.allow(drained, start.Add(59*time.Second))
	limiter.allow(drained, start.Add(59*time.Second))

	// A minute after the first sweep, this call sweeps again.
	limiter.allow("198.51.100.3", start.Add(time.Minute))

	if _, ok := limiter.byAddr[refilled]; ok {
		t.Error("address with a full bucket was kept")
	}

	if _, ok := limiter.byAddr[drained]; !ok {
		t.Error("address short of tokens was dropped")
	}
}