1package ratelimit
2
3import (
4 "context"
5 "sync"
6 "testing"
7 "time"
8
9 "github.com/bnema/zerowrap"
10 "github.com/stretchr/testify/assert"
11 "github.com/stretchr/testify/require"
12)
13
14func testLogger() zerowrap.Logger {
15 return zerowrap.Default()
16}
17
18func TestMemoryStore_Allow_WithinLimit(t *testing.T) {
19 store := NewMemoryStore(10, 10, testLogger()) // 10 RPS, burst 10
20 ctx := context.Background()
21
22 // Should allow requests up to burst
23 for i := 0; i < 10; i++ {
24 assert.True(t, store.Allow(ctx, "test"), "request %d should be allowed", i+1)
25 }
26}
27
28func TestMemoryStore_Allow_ExceedsLimit(t *testing.T) {
29 store := NewMemoryStore(1, 1, testLogger()) // 1 RPS, burst 1
30 ctx := context.Background()
31
32 // First request should succeed (consumes the burst)
33 assert.True(t, store.Allow(ctx, "test"), "first request should be allowed")
34
35 // Second request should be rate limited
36 assert.False(t, store.Allow(ctx, "test"), "second request should be rate limited")
37}
38
39func TestMemoryStore_Allow_BurstBehavior(t *testing.T) {
40 store := NewMemoryStore(10, 5, testLogger()) // 10 RPS, burst 5
41 ctx := context.Background()
42
43 // Should allow burst of 5
44 for i := 0; i < 5; i++ {
45 assert.True(t, store.Allow(ctx, "test"), "burst request %d should be allowed", i+1)
46 }
47
48 // 6th request should be rate limited (burst exhausted)
49 assert.False(t, store.Allow(ctx, "test"), "request exceeding burst should be rate limited")
50
51 // Wait for token replenishment
52 time.Sleep(200 * time.Millisecond)
53
54 // Should be able to make at least one request after waiting
55 assert.True(t, store.Allow(ctx, "test"), "request after waiting should be allowed")
56}
57
58func TestMemoryStore_Allow_IndependentKeys(t *testing.T) {
59 store := NewMemoryStore(1, 1, testLogger()) // 1 RPS, burst 1
60 ctx := context.Background()
61
62 // First request for key1 should succeed
63 assert.True(t, store.Allow(ctx, "key1"), "first request for key1 should be allowed")
64
65 // Key1 is now rate limited
66 assert.False(t, store.Allow(ctx, "key1"), "second request for key1 should be rate limited")
67
68 // Key2 should have its own independent limit
69 assert.True(t, store.Allow(ctx, "key2"), "first request for key2 should be allowed")
70
71 // Key2 is now rate limited
72 assert.False(t, store.Allow(ctx, "key2"), "second request for key2 should be rate limited")
73}
74
75func TestMemoryStore_AllowN_Basic(t *testing.T) {
76 store := NewMemoryStore(10, 10, testLogger()) // 10 RPS, burst 10
77 ctx := context.Background()
78
79 // Should allow 5 tokens at once
80 assert.True(t, store.AllowN(ctx, "test", 5), "AllowN(5) should succeed with burst of 10")
81
82 // Should allow another 5
83 assert.True(t, store.AllowN(ctx, "test", 5), "AllowN(5) should succeed with 5 remaining")
84
85 // Should not allow more
86 assert.False(t, store.AllowN(ctx, "test", 1), "AllowN(1) should fail with 0 remaining")
87}
88
89func TestMemoryStore_AllowN_ExceedsBurst(t *testing.T) {
90 store := NewMemoryStore(10, 5, testLogger()) // 10 RPS, burst 5
91 ctx := context.Background()
92
93 // Should not allow more than burst
94 assert.False(t, store.AllowN(ctx, "test", 10), "AllowN(10) should fail with burst of 5")
95
96 // But smaller request should work
97 assert.True(t, store.AllowN(ctx, "test", 3), "AllowN(3) should succeed with burst of 5")
98}
99
100func TestMemoryStore_Allow_ConcurrentAccess(t *testing.T) {
101 store := NewMemoryStore(1000, 100, testLogger()) // High RPS for concurrent testing
102 ctx := context.Background()
103
104 var wg sync.WaitGroup
105 results := make(chan bool, 200)
106
107 // Run 20 goroutines, each making 10 requests
108 for i := 0; i < 20; i++ {
109 wg.Add(1)
110 go func() {
111 defer wg.Done()
112 for j := 0; j < 10; j++ {
113 results <- store.Allow(ctx, "concurrent")
114 }
115 }()
116 }
117
118 wg.Wait()
119 close(results)
120
121 // Count successful requests
122 allowed := 0
123 for result := range results {
124 if result {
125 allowed++
126 }
127 }
128
129 // At least burst number should be allowed
130 require.GreaterOrEqual(t, allowed, 100, "at least burst number of requests should be allowed")
131 // And we shouldn't exceed total requests
132 require.LessOrEqual(t, allowed, 200, "should not exceed total requests")
133}
134
135func TestMemoryStore_IPKeyFormat(t *testing.T) {
136 store := NewMemoryStore(1, 1, testLogger())
137 ctx := context.Background()
138
139 // Test typical IP key formats
140 assert.True(t, store.Allow(ctx, "ip:192.168.1.1"), "should allow first request for IP")
141 assert.False(t, store.Allow(ctx, "ip:192.168.1.1"), "should rate limit same IP")
142 assert.True(t, store.Allow(ctx, "ip:192.168.1.2"), "should allow first request for different IP")
143}
144
145func TestMemoryStore_GlobalKey(t *testing.T) {
146 store := NewMemoryStore(1, 1, testLogger())
147 ctx := context.Background()
148
149 // Test global key
150 assert.True(t, store.Allow(ctx, "global"), "should allow first global request")
151 assert.False(t, store.Allow(ctx, "global"), "should rate limit global requests")
152}