package middleware

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

func TestRateLimit(t *testing.T) {
	l := NewKeyedLimiter(t.Context(), 1, 2, 0)
	next := http.HandlerFunc(func(http.ResponseWriter, *http.Request) {
	})
	h := RateLimit(l, nil)(next)

	var last int
	for i := 0; i < 4; i++ {
		rec := httptest.NewRecorder()
		h.ServeHTTP(rec, httptest.NewRequest("GET", "/", nil))
		last = rec.Code
		if last == http.StatusTooManyRequests &&
			rec.Header().Get("Retry-After") == "" {
			t.Fatal("missing Retry-After")
		}
	}

	if last != http.StatusTooManyRequests {
		t.Fatalf("want 429, got %d", last)
	}
}