package ratelimit

import (
	"testing"
	"time"
)

// TestBurstThenDeny is the property the relay depends on: a key gets exactly
// burst writes, and the state has to survive between calls.
func TestBurstThenDeny(t *testing.T) {
	l := New(0.000001, 3)
	key := []byte("1.2.3.4")
	for i := int32(0); i < 3; i++ {
		if !l.Allow(key) {
			t.Fatalf("token %d should be allowed", i)
		}
	}
	if l.Allow(key) {
		t.Fatal("burst is exhausted, the next write must be denied")
	}
	if len(l.buckets) != 1 {
		t.Fatalf("one bucket expected, got %d", len(l.buckets))
	}
}

func TestDistinctKeys(t *testing.T) {
	l := New(0.000001, 1)
	if !l.Allow([]byte("a")) {
		t.Fatal("first key should be allowed")
	}
	if l.Allow([]byte("a")) {
		t.Fatal("first key is spent")
	}
	if !l.Allow([]byte("b")) {
		t.Fatal("a second key has its own burst")
	}
}

func TestRefill(t *testing.T) {
	// 10 tokens/second: one token needs 100ms, so the immediate second call
	// cannot have refilled.
	l := New(10.0, 1)
	key := []byte("k")
	if !l.Allow(key) {
		t.Fatal("first allow")
	}
	if l.Allow(key) {
		t.Fatal("no token should be left")
	}
	time.Sleep(200 * time.Millisecond)
	if !l.Allow(key) {
		t.Fatal("the bucket should have refilled")
	}
}

// TestRefillIsCapped checks that a long idle period does not accumulate more
// than burst tokens.
func TestRefillIsCapped(t *testing.T) {
	l := New(10.0, 2)
	key := []byte("k")
	time.Sleep(300 * time.Millisecond)
	for i := int32(0); i < 2; i++ {
		if !l.Allow(key) {
			t.Fatalf("token %d should be allowed", i)
		}
	}
	if l.Allow(key) {
		t.Fatal("burst is 2, a third token must not be granted")
	}
}

func TestCleanup(t *testing.T) {
	l := New(1.0, 1)
	if !l.Allow([]byte("old")) {
		t.Fatal("allow")
	}
	if len(l.buckets) != 1 {
		t.Fatal("bucket should be recorded")
	}
	l.Cleanup(time.Duration(0))
	if len(l.buckets) != 0 {
		t.Fatal("Cleanup(0) should drop every bucket")
	}
	// The limiter still works after its state was dropped.
	if !l.Allow([]byte("new")) {
		t.Fatal("allow after cleanup")
	}
}