learning-go-for-sre-interview/03-rate-limiter/main_test.go

package main

import (
	"sync"
	"sync/atomic"
	"testing"
	"time"
)

func TestBurstAllowedThenDenied(t *testing.T) {
	l := NewLimiter(1, 3) // slow refill so it doesn't interfere
	for i := 0; i < 3; i++ {
		if !l.Allow() {
			t.Fatalf("request %d within burst should be allowed", i)
		}
	}
	if l.Allow() {
		t.Fatal("4th immediate request should be denied (bucket empty)")
	}
}

func TestRefillOverTime(t *testing.T) {
	l := NewLimiter(20, 1) // 20 tokens/sec => 1 token per 50ms
	if !l.Allow() {
		t.Fatal("initial token missing")
	}
	if l.Allow() {
		t.Fatal("bucket should be empty immediately after")
	}
	time.Sleep(120 * time.Millisecond) // ~2.4 tokens accrued, capped at 1
	if !l.Allow() {
		t.Fatal("token should have refilled after wait")
	}
}

func TestBurstIsCapped(t *testing.T) {
	l := NewLimiter(1000, 2)
	time.Sleep(50 * time.Millisecond) // would accrue ~50 tokens uncapped
	allowed := 0
	for i := 0; i < 10; i++ {
		if l.Allow() {
			allowed++
		}
	}
	if allowed > 2 {
		t.Fatalf("bucket exceeded burst: %d immediate allows (cap 2) — idle time built infinite credit", allowed)
	}
}

func TestConcurrentNoOverAdmission(t *testing.T) {
	l := NewLimiter(0.0001, 100) // effectively no refill during the test
	var allowed atomic.Int64
	var wg sync.WaitGroup
	for i := 0; i < 50; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			for j := 0; j < 10; j++ {
				if l.Allow() {
					allowed.Add(1)
				}
			}
		}()
	}
	wg.Wait()
	if got := allowed.Load(); got > 100 {
		t.Fatalf("over-admission under concurrency: %d > 100 tokens spent", got)
	} else if got < 100 {
		t.Fatalf("under-admission: only %d of 100 available tokens spent", got)
	}
}