Skip to content

Commit e21cc36

Browse files
committed
Add redis tests and fix bugs
1 parent 98dbfac commit e21cc36

File tree

7 files changed

+147
-21
lines changed

7 files changed

+147
-21
lines changed

.gitattributes

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
* text eol=lf

go.mod

Lines changed: 16 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -3,8 +3,19 @@ module github.com/Pavan-Rana/rate-limiter
33
go 1.22
44

55
require (
6-
github.com/prometheus/client_golang v1.19.0
7-
github.com/redis/go-redis/v9 v9.5.1
8-
google.golang.org/grpc v1.63.2
9-
google.golang.org/protobuf v1.34.1
10-
)
6+
github.com/prometheus/client_golang v1.19.0
7+
github.com/redis/go-redis/v9 v9.5.1
8+
google.golang.org/grpc v1.63.2
9+
google.golang.org/protobuf v1.34.1
10+
)
11+
12+
require (
13+
github.com/alicebob/miniredis/v2 v2.37.0 // indirect
14+
github.com/cespare/xxhash/v2 v2.2.0 // indirect
15+
github.com/davecgh/go-spew v1.1.1 // indirect
16+
github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f // indirect
17+
github.com/pmezard/go-difflib v1.0.0 // indirect
18+
github.com/stretchr/testify v1.11.1 // indirect
19+
github.com/yuin/gopher-lua v1.1.1 // indirect
20+
gopkg.in/yaml.v3 v3.0.1 // indirect
21+
)

go.sum

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,27 @@
1+
github.com/alicebob/miniredis/v2 v2.37.0 h1:RheObYW32G1aiJIj81XVt78ZHJpHonHLHW7OLIshq68=
2+
github.com/alicebob/miniredis/v2 v2.37.0/go.mod h1:TcL7YfarKPGDAthEtl5NBeHZfeUQj6OXMm/+iu5cLMM=
3+
github.com/cespare/xxhash/v2 v2.2.0 h1:DC2CZ1Ep5Y4k3ZQ899DldepgrayRUGE6BBZ/cd9Cj44=
4+
github.com/cespare/xxhash/v2 v2.2.0/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs=
5+
github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
6+
github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
7+
github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f h1:lO4WD4F/rVNCu3HqELle0jiPLLBs70cWOduZpkS1E78=
8+
github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f/go.mod h1:cuUVRXasLTGF7a8hSLbxyZXjz+1KgoB3wDUb6vlszIc=
9+
github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
10+
github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
111
github.com/prometheus/client_golang v1.19.0 h1:ygXvpU1AoN1MhdzckN+PyD9QJOSD4x7kmXYlnfbA6JU=
212
github.com/prometheus/client_golang v1.19.0/go.mod h1:ZRM9uEAypZakd+q/x7+gmsvXdURP+DABIEIjnmDdp+k=
313
github.com/redis/go-redis/v9 v9.5.1 h1:H1X4D3yHPaYrkL5X06Wh6xNVM/pX0Ft4RV0vMGvLBh8=
414
github.com/redis/go-redis/v9 v9.5.1/go.mod h1:hdY0cQFCN4fnSYT6TkisLufl/4W5UIXyv0b/CLO2V2M=
15+
github.com/stretchr/testify v1.11.1 h1:7s2iGBzp5EwR7/aIZr8ao5+dra3wiQyKjjFuvgVKu7U=
16+
github.com/stretchr/testify v1.11.1/go.mod h1:wZwfW3scLgRK+23gO65QZefKpKQRnfz6sD981Nm4B6U=
17+
github.com/yuin/gopher-lua v1.1.1 h1:kYKnWBjvbNP4XLT3+bPEwAXJx262OhaHDWDVOPjL46M=
18+
github.com/yuin/gopher-lua v1.1.1/go.mod h1:GBR0iDaNXjAgGg9zfCvksxSRnQx76gclCIb7kdAd1Pw=
519
google.golang.org/grpc v1.63.2 h1:MUeiw1B2maTVZthpU5xvASfTh3LDbxHd6IJ6QQVU+xM=
620
google.golang.org/grpc v1.63.2/go.mod h1:WAX/8DgncnokcFUldAxq7GeB5DXHDbMF+lLvDomNkRA=
721
google.golang.org/protobuf v1.33.0 h1:uNO2rsAINq/JlFpSdYEKIZ0uKD/R9cpdv0T+yoGwGmI=
822
google.golang.org/protobuf v1.33.0/go.mod h1:c6P6GXX6sHbq/GpV6MGZEdwhWPcYBgnhAHhKbcUYpos=
923
google.golang.org/protobuf v1.34.1 h1:9ddQBjfCyZPOHPUiPxpYESBLc+T8P3E+Vo4IbKZgFWg=
1024
google.golang.org/protobuf v1.34.1/go.mod h1:c6P6GXX6sHbq/GpV6MGZEdwhWPcYBgnhAHhKbcUYpos=
25+
gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
26+
gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=
27+
gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=

internal/store/lua_scripts.go

Lines changed: 14 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -2,15 +2,12 @@ package store
22

33
import (
44
"context"
5-
"embed"
5+
_ "embed"
6+
"os"
67

78
"github.com/redis/go-redis/v9"
89
)
910

10-
var slidingWindowScript string
11-
12-
var tokenBucketScript string
13-
1411
// LuaScripts holds SHA digests for pre-loaded Lua scripts.
1512
// Using EVALSHA (vs EVAL) avoids re-parsing the script on every request.
1613
type LuaScripts struct {
@@ -21,12 +18,21 @@ type LuaScripts struct {
2118
// LoadScripts pre-loads Lua scripts into Redis at startup.
2219
// Subsequent calls use EVALSHA with the digest - faster and avoids redundant parsing.
2320
func LoadScripts(ctx context.Context, client *redis.Client) (*LuaScripts, error) {
24-
swSHA, err := client.ScriptLoad(ctx, slidingWindowScript).Result()
21+
sw, err := os.ReadFile("../../scripts/lua/sliding_window.lua")
22+
if err != nil {
23+
return nil, err
24+
}
25+
tb, err := os.ReadFile("../../scripts/lua/token_bucket.lua")
2526
if err != nil {
2627
return nil, err
2728
}
2829

29-
tbSHA, err := client.ScriptLoad(ctx, tokenBucketScript).Result()
30+
swSHA, err := client.ScriptLoad(ctx, string(sw)).Result()
31+
if err != nil {
32+
return nil, err
33+
}
34+
35+
tbSHA, err := client.ScriptLoad(ctx, string(tb)).Result()
3036
if err != nil {
3137
return nil, err
3238
}
@@ -35,4 +41,4 @@ func LoadScripts(ctx context.Context, client *redis.Client) (*LuaScripts, error)
3541
SlidingWindowSHA: swSHA,
3642
TokenBucketSHA: tbSHA,
3743
}, nil
38-
}
44+
}

internal/store/redis.go

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -4,13 +4,13 @@ import (
44
"context"
55
"time"
66

7+
"github.com/Pavan-Rana/rate-limiter/internal/limiter"
78
"github.com/redis/go-redis/v9"
8-
"github.com/Pavan-Rana/redis-rate-limiter/internal/limiter"
99
)
1010

1111
type RedisStore struct {
12-
client *redis.Client
13-
scripts *LuaScripts
12+
client *redis.Client
13+
scripts *LuaScripts
1414
}
1515

1616
func NewRedisStore(addr string) (*RedisStore, error) {
@@ -36,7 +36,7 @@ func (r *RedisStore) AllowRequest(ctx context.Context, key string, policy limite
3636
windowMs := policy.Window.Milliseconds()
3737

3838
result, err := r.client.EvalSha(ctx, r.scripts.SlidingWindowSHA, []string{key},
39-
nowMs, windowMs, policy.Limit
39+
nowMs, windowMs, policy.Limit,
4040
).Int()
4141

4242
if err != nil {
@@ -48,4 +48,4 @@ func (r *RedisStore) AllowRequest(ctx context.Context, key string, policy limite
4848

4949
func (r *RedisStore) Ping(ctx context.Context) error {
5050
return r.client.Ping(ctx).Err()
51-
}
51+
}

internal/store/redis_test.go

Lines changed: 92 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1 +1,92 @@
1-
// internal/store/redis_test.go
1+
package store
2+
3+
import (
4+
"context"
5+
"testing"
6+
"time"
7+
8+
"github.com/alicebob/miniredis/v2"
9+
"github.com/stretchr/testify/assert"
10+
"github.com/stretchr/testify/require"
11+
12+
"github.com/Pavan-Rana/rate-limiter/internal/limiter"
13+
)
14+
15+
func newTestStore(t *testing.T) (*RedisStore, *miniredis.Miniredis) {
16+
t.Helper()
17+
mr := miniredis.RunT(t)
18+
store, err := NewRedisStore(mr.Addr())
19+
require.NoError(t, err)
20+
return store, mr
21+
}
22+
23+
func policy(limit int, window time.Duration) limiter.Policy {
24+
return limiter.Policy{Limit: limit, Window: window}
25+
}
26+
27+
func TestNewRedisStore_UnreachableServer(t *testing.T) {
28+
store, err := NewRedisStore("127.0.0.1:1")
29+
assert.Error(t, err)
30+
assert.Nil(t, store)
31+
}
32+
33+
func TestPing_AfterServerDown(t *testing.T) {
34+
store, mr := newTestStore(t)
35+
mr.Close()
36+
assert.Error(t, store.Ping(context.Background()))
37+
}
38+
39+
func TestAllowRequest_WithinAndBeyondLimit(t *testing.T) {
40+
store, _ := newTestStore(t)
41+
p := policy(3, time.Minute)
42+
ctx := context.Background()
43+
44+
for i := 0; i < p.Limit; i++ {
45+
allowed, err := store.AllowRequest(ctx, "client:limit", p)
46+
require.NoError(t, err)
47+
assert.True(t, allowed, "request %d should be allowed", i+1)
48+
}
49+
50+
allowed, err := store.AllowRequest(ctx, "client:limit", p)
51+
require.NoError(t, err)
52+
assert.False(t, allowed, "request beyond limit should be denied")
53+
}
54+
55+
func TestAllowRequest_WindowExpiry_ResetsCount(t *testing.T) {
56+
store, mr := newTestStore(t)
57+
window := 100 * time.Millisecond
58+
p := policy(1, window)
59+
ctx := context.Background()
60+
61+
_, err := store.AllowRequest(ctx, "client:expiry", p)
62+
require.NoError(t, err)
63+
64+
mr.FastForward(window + 10*time.Millisecond)
65+
66+
allowed, err := store.AllowRequest(ctx, "client:expiry", p)
67+
require.NoError(t, err)
68+
assert.True(t, allowed, "should be allowed after window expiry")
69+
}
70+
71+
func TestAllowRequest_DifferentKeysAreIsolated(t *testing.T) {
72+
store, _ := newTestStore(t)
73+
p := policy(1, time.Minute)
74+
ctx := context.Background()
75+
76+
_, err := store.AllowRequest(ctx, "client:A", p)
77+
require.NoError(t, err)
78+
79+
allowed, err := store.AllowRequest(ctx, "client:B", p)
80+
require.NoError(t, err)
81+
assert.True(t, allowed, "different keys should not share rate-limit state")
82+
}
83+
84+
func TestAllowRequest_CancelledContext(t *testing.T) {
85+
store, _ := newTestStore(t)
86+
87+
ctx, cancel := context.WithCancel(context.Background())
88+
cancel()
89+
90+
_, err := store.AllowRequest(ctx, "client:ctx", policy(10, time.Minute))
91+
assert.Error(t, err)
92+
}

scripts/lua/sliding_window.lua

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -39,9 +39,9 @@ if count < limit then
3939
local member = tostring(now) .. '-' .. tostring(redis.call('INCR', key .. ':seq'))
4040
redis.call('ZADD', key, now, member)
4141

42-
redis.call('PEXPIRE', key, window_ms + 1000)
42+
redis.call('PEXPIRE', key, window_ms)
4343

4444
return 1
45-
else:
45+
else
4646
return 0
4747
end

0 commit comments

Comments
 (0)