x/limiter/quota/limiter_test.go

package quota

import (
	"encoding/json"
	"os"
	"path/filepath"
	"testing"
	"time"
)

func newTestLimiter(t *testing.T, opts Options) *Limiter {
	t.Helper()
	l := NewLimiter("test", opts)
	t.Cleanup(func() { l.Close() })
	return l
}

func TestLimiterBlocksAtLimit(t *testing.T) {
	l := newTestLimiter(t, Options{Limit: 100})

	l.AddIn(60)
	if l.Blocked() {
		t.Fatalf("should not block at 60/100")
	}
	l.AddOut(40)
	if !l.Blocked() {
		t.Fatalf("should block at 100/100")
	}
	if s := l.Snapshot(); !s.Blocked || s.Used != 100 {
		t.Fatalf("snapshot=%+v, want blocked used=100", s)
	}
}

func TestLimiterUnblockOnOverwrite(t *testing.T) {
	l := newTestLimiter(t, Options{Limit: 100})
	l.AddIn(150)
	if !l.Blocked() {
		t.Fatalf("expected blocked")
	}

	zero := uint64(0)
	l.Update(Update{Used: &zero})

	if l.Blocked() {
		t.Fatalf("expected unblocked after reset")
	}
	if got := l.Snapshot().Used; got != 0 {
		t.Fatalf("used=%d, want 0", got)
	}
}

func TestLimiterDirection(t *testing.T) {
	l := newTestLimiter(t, Options{Limit: 100, Direction: DirectionIn})

	l.AddOut(500)
	if u := l.Snapshot().Used; u != 0 {
		t.Fatalf("outbound counted under DirectionIn: used=%d", u)
	}
	l.AddIn(100)
	if !l.Blocked() {
		t.Fatalf("expected blocked after inbound reaches limit")
	}
}

func TestLimiterWindowNotStarted(t *testing.T) {
	l := newTestLimiter(t, Options{Limit: 100, StartsAt: time.Now().Add(time.Hour)})

	l.AddIn(500)
	s := l.Snapshot()
	if s.Used != 0 {
		t.Fatalf("counted before startsAt: used=%d", s.Used)
	}
	if s.Active {
		t.Fatalf("should be inactive before startsAt")
	}
	if l.Blocked() {
		t.Fatalf("should not block before startsAt")
	}
}

func TestLimiterFailOpenAfterExpiry(t *testing.T) {
	l := newTestLimiter(t, Options{Limit: 100, ExpiresAt: time.Now().Add(-time.Hour)})

	l.AddIn(500)
	s := l.Snapshot()
	if s.Used != 0 {
		t.Fatalf("counted after expiry: used=%d", s.Used)
	}
	if !s.Expired {
		t.Fatalf("snapshot should report expired")
	}
	if l.Blocked() {
		t.Fatalf("should be fail-open (unblocked) after expiry")
	}
}

func TestLimiterPersistsToDisk(t *testing.T) {
	path := filepath.Join(t.TempDir(), "quota.json")
	store := NewFileStore(path)

	l := NewLimiter("svc", Options{Limit: 1000, Store: store})
	l.AddIn(300)
	l.AddOut(200)
	l.Close()

	b, err := os.ReadFile(path)
	if err != nil {
		t.Fatalf("read store file: %v", err)
	}
	var data map[string]Record
	if err := json.Unmarshal(b, &data); err != nil {
		t.Fatalf("unmarshal: %v", err)
	}
	if rec := data["svc"]; rec.Used != 500 {
		t.Fatalf("persisted used=%d, want 500", rec.Used)
	}
}

func TestLimiterRestoresFromDisk(t *testing.T) {
	path := filepath.Join(t.TempDir(), "quota.json")
	if err := os.WriteFile(path, []byte(`{"svc":{"used":777,"limit":1000}}`), 0o644); err != nil {
		t.Fatal(err)
	}

	store := NewFileStore(path)
	l := NewLimiter("svc", Options{Limit: 1000, Store: store})
	t.Cleanup(func() { l.Close() })

	if got := l.Snapshot().Used; got != 777 {
		t.Fatalf("restored used=%d, want 777", got)
	}
}

func TestWaitChanNotifiesOnReset(t *testing.T) {
	l := newTestLimiter(t, Options{Limit: 100})
	l.AddIn(200)
	if !l.Blocked() {
		t.Fatal("expected blocked")
	}

	ch := l.WaitChan()
	zero := uint64(0)
	l.Update(Update{Used: &zero})

	select {
	case <-ch:
	case <-time.After(time.Second):
		t.Fatal("WaitChan not notified after reset")
	}
	if l.Blocked() {
		t.Fatal("expected unblocked after reset")
	}
}

func TestClosedLimiterIsInert(t *testing.T) {
	l := NewLimiter("c", Options{Limit: 100})
	l.AddIn(200)
	if !l.Blocked() {
		t.Fatal("expected blocked")
	}

	ch := l.WaitChan()
	l.Close()

	select {
	case <-ch:
	case <-time.After(time.Second):
		t.Fatal("WaitChan not notified on close")
	}
	if l.Blocked() {
		t.Fatal("closed limiter must be inert (not blocking)")
	}
	l.AddIn(1000)
	if u := l.Snapshot().Used; u != 200 {
		t.Fatalf("closed limiter counted: used=%d, want 200", u)
	}
}

func TestLimiterResetsOnWindowChange(t *testing.T) {
	path := filepath.Join(t.TempDir(), "quota.json")
	store := NewFileStore(path)

	t1 := time.Now().Add(time.Hour)
	l1 := NewLimiter("svc", Options{Limit: 1000, ExpiresAt: t1, Store: store})
	l1.AddIn(500)
	l1.Close()

	t2 := time.Now().Add(2 * time.Hour)
	l2 := NewLimiter("svc", Options{Limit: 1000, ExpiresAt: t2, Store: store})
	t.Cleanup(func() { l2.Close() })

	if u := l2.Snapshot().Used; u != 0 {
		t.Fatalf("counter not reset on window change: used=%d, want 0", u)
	}
}