package ratelimit

import (
	"net/http"
	"testing"
	"time"
)

func TestTokenBucket(t *testing.T) {
	l := New(1, 2) // 1 token/sec, burst 2
	now := time.Unix(1000, 0)
	l.now = func() time.Time { return now }

	// Burst of 2 is allowed, the third in the same instant is denied.
	if !l.Allow("a") || !l.Allow("a") {
		t.Fatal("burst of 2 should be allowed")
	}
	if l.Allow("a") {
		t.Fatal("third request in the same instant must be denied")
	}
	// A different key has its own bucket.
	if !l.Allow("b") {
		t.Fatal("a separate key must not share key a's bucket")
	}
	// After one second, one token refills.
	now = now.Add(time.Second)
	if !l.Allow("a") {
		t.Fatal("one token should refill after a second")
	}
	if l.Allow("a") {
		t.Fatal("only one token refilled, the next must be denied")
	}
}

func TestSweepDropsIdleFullBuckets(t *testing.T) {
	l := New(1, 1)
	now := time.Unix(1000, 0)
	l.now = func() time.Time { return now }
	l.Allow("a") // creates the bucket, now empty
	now = now.Add(10 * time.Second)
	l.Allow("a") // refills to full and consumes... back below full, so not swept yet
	now = now.Add(10 * time.Second)
	// "a" has refilled to full (capacity 1) by now; sweep should drop it.
	l.sweep(5 * time.Second)
	l.mu.Lock()
	_, present := l.buckets["a"]
	l.mu.Unlock()
	if present {
		t.Fatal("an idle, full bucket should be swept")
	}
}

func TestClientIP(t *testing.T) {
	r := &http.Request{RemoteAddr: "203.0.113.7:54321", Header: http.Header{}}
	if got := ClientIP(r, ""); got != "203.0.113.7" {
		t.Fatalf("TCP peer: got %q want 203.0.113.7", got)
	}
	r.Header.Set("X-Real-IP", "198.51.100.9")
	if got := ClientIP(r, "X-Real-IP"); got != "198.51.100.9" {
		t.Fatalf("trusted header: got %q want 198.51.100.9", got)
	}
	// A spoofed header is ignored when no trusted header is configured.
	if got := ClientIP(r, ""); got != "203.0.113.7" {
		t.Fatalf("untrusted: got %q want the TCP peer", got)
	}
	// A proxy chain takes the first entry.
	r.Header.Set("X-Real-IP", "198.51.100.9, 10.0.0.1")
	if got := ClientIP(r, "X-Real-IP"); got != "198.51.100.9" {
		t.Fatalf("proxy chain: got %q want the first entry", got)
	}
}