ratelimit_test.gno

  1package ratelimit
  2
  3import (
  4	"testing"
  5
  6	"gno.land/p/nt/uassert/v0"
  7)
  8
  9func TestRefill(t *testing.T) {
 10	// Start with 1 token, 4 ticks elapse at 0.5 tokens/tick, cap 5.
 11	// 1 + 4*0.5 = 3.
 12	got := refill(1.0, 10, 14, 0.5, 5.0)
 13	uassert.True(t, got == 3.0, "expected 3 tokens after refill")
 14
 15	// No time elapsed -> unchanged (but clamped to cap).
 16	got = refill(2.0, 10, 10, 0.5, 5.0)
 17	uassert.True(t, got == 2.0, "no elapsed ticks should not add tokens")
 18
 19	// Going backwards (now < last) never adds tokens.
 20	got = refill(2.0, 10, 5, 0.5, 5.0)
 21	uassert.True(t, got == 2.0, "negative elapsed should not add tokens")
 22}
 23
 24func TestRefillCap(t *testing.T) {
 25	// Long idle period must saturate at burst capacity.
 26	got := refill(0.0, 0, 1000, 0.5, 5.0)
 27	uassert.True(t, got == 5.0, "refill must cap at burst")
 28
 29	// Already-full bucket stays full.
 30	got = refill(5.0, 0, 100, 0.5, 5.0)
 31	uassert.True(t, got == 5.0, "full bucket stays capped")
 32}
 33
 34func TestTokensToInt(t *testing.T) {
 35	uassert.Equal(t, 3, tokensToInt(3.9))
 36	uassert.Equal(t, 0, tokensToInt(0.4))
 37	uassert.Equal(t, 0, tokensToInt(-2.0))
 38	uassert.Equal(t, 5, tokensToInt(5.0))
 39}
 40
 41func TestFmin(t *testing.T) {
 42	uassert.True(t, fmin(1.0, 2.0) == 1.0, "fmin picks smaller")
 43	uassert.True(t, fmin(4.0, 2.5) == 2.5, "fmin picks smaller")
 44	uassert.True(t, fmin(3.0, 3.0) == 3.0, "fmin on equals")
 45}
 46
 47func TestNewClampsConfig(t *testing.T) {
 48	l := New(-1, 0) // rate < 0 -> 0, burst < 1 -> 1
 49	r, b := l.Config()
 50	uassert.True(t, r == 0.0, "rate clamped to 0")
 51	uassert.True(t, b == 1.0, "burst clamped to 1")
 52}
 53
 54func TestAllowConsumesAndDenies(t *testing.T) {
 55	l := New(0.5, 5.0)
 56	const k = "g1caller"
 57	// A fresh key starts full (burst=5): first 5 Allow succeed, 6th fails.
 58	for i := 0; i < 5; i++ {
 59		uassert.True(t, l.Allow(k, 0), "burst token should be granted")
 60	}
 61	uassert.False(t, l.Allow(k, 0), "6th call at same tick must be denied")
 62
 63	// After 4 ticks at 0.5/tick, 2 whole tokens refilled -> 2 more allowed.
 64	uassert.True(t, l.Allow(k, 4), "refilled token granted")
 65	uassert.True(t, l.Allow(k, 4), "refilled token granted")
 66	uassert.False(t, l.Allow(k, 4), "empty again")
 67}
 68
 69func TestAllowNAtomic(t *testing.T) {
 70	l := New(1, 5.0)
 71	const k = "g1caller"
 72	uassert.False(t, l.AllowN(k, 0, 6), "cannot take more than burst")
 73	uassert.Equal(t, 5, l.Tokens(k, 0), "denied AllowN consumes nothing")
 74	uassert.True(t, l.AllowN(k, 0, 3), "3 of 5 taken")
 75	uassert.Equal(t, 2, l.Tokens(k, 0), "2 remain")
 76}
 77
 78func TestTokensReadOnly(t *testing.T) {
 79	l := New(0.5, 5.0)
 80	const k = "g1caller"
 81	// Unseen key reports a full bucket without creating it.
 82	uassert.Equal(t, 5, l.Tokens(k, 100), "unseen key is full")
 83	uassert.Equal(t, 0, l.Len(), "Tokens must not create a bucket")
 84	// Reading does not consume.
 85	uassert.Equal(t, 5, l.Tokens(k, 200), "still full on second read")
 86}
 87
 88func TestPerKeyIsolation(t *testing.T) {
 89	l := New(0.5, 5.0)
 90	l.Allow("g1a", 0)
 91	uassert.Equal(t, 4, l.Tokens("g1a", 0), "a spent one")
 92	uassert.Equal(t, 5, l.Tokens("g1b", 0), "b untouched")
 93	uassert.Equal(t, 1, l.Len(), "only a exists")
 94}
 95
 96func TestIterate(t *testing.T) {
 97	l := New(0.5, 5.0)
 98	l.Allow("g1a", 0)
 99	l.Allow("g1b", 0)
100	seen := 0
101	stopped := l.Iterate(0, func(key string, tokens int, last int64) bool {
102		seen++
103		return false
104	})
105	uassert.Equal(t, 2, seen, "iterate visits every key")
106	uassert.False(t, stopped, "not stopped early")
107
108	// Early stop.
109	count := 0
110	stopped = l.Iterate(0, func(key string, tokens int, last int64) bool {
111		count++
112		return true
113	})
114	uassert.Equal(t, 1, count, "stops after first")
115	uassert.True(t, stopped, "reported as stopped")
116}