Skip to content

Commit bebd6fa

Browse files
fforbeckfforbeck
committed
testing rate-limit logic
1 parent 8646706 commit bebd6fa

File tree

1 file changed

+151
-13
lines changed

1 file changed

+151
-13
lines changed

‎test/middlewares/rate-limiter.spec.js‎

Lines changed: 151 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,5 @@
1-
import { withRateLimits } from '../../src/middleware.js'
21
import { jest } from '@jest/globals'
3-
import { RATE_LIMIT_EXCEEDED } from '../../src/constants.js'
2+
import { withRateLimits } from '../../src/middleware.js'
43
import { HttpError } from '@web3-storage/gateway-lib/util'
54

65
describe('withRateLimits', () => {
@@ -13,52 +12,191 @@ describe('withRateLimits', () => {
1312
rateLimiter = {
1413
limit: jest.fn()
1514
}
16-
env = { RATE_LIMITER: rateLimiter, ACCOUNTING_SERVICE_URL: 'http://example.com' }
15+
env = {
16+
RATE_LIMITER: rateLimiter,
17+
ACCOUNTING_SERVICE_URL: 'http://example.com',
18+
AUTH_TOKEN_METADATA: {
19+
get: jest.fn(),
20+
put: jest.fn()
21+
}
22+
}
1723
next = jest.fn()
1824
handler = jest.fn()
1925
})
2026

21-
it('should call next if rate limit is not exceeded', async () => {
27+
afterEach(() => {
28+
jest.restoreAllMocks()
29+
})
30+
31+
it('should call next if no auth token and rate limit is not exceeded', async () => {
32+
rateLimiter.limit.mockResolvedValue({ success: true })
33+
34+
const request = {
35+
headers: {
36+
get: jest.fn()
37+
}
38+
}
39+
const ctx = { dataCid: 'test-cid' }
40+
const wrappedHandler = withRateLimits(handler)
41+
42+
await wrappedHandler(request, env, ctx)
43+
44+
expect(rateLimiter.limit).toHaveBeenCalledTimes(1)
45+
expect(rateLimiter.limit).toHaveBeenCalledWith({ key: ctx.dataCid.toString() })
46+
expect(handler).toHaveBeenCalledTimes(1)
47+
expect(handler).toHaveBeenCalledWith(request, env, ctx)
48+
})
49+
50+
it('should throw an error if no auth token and rate limit is exceeded', async () => {
51+
rateLimiter.limit.mockResolvedValue({ success: false })
52+
53+
const request = {
54+
headers: {
55+
get: jest.fn()
56+
}
57+
}
58+
const ctx = { dataCid: 'test-cid' }
59+
const wrappedHandler = withRateLimits(handler)
60+
61+
try {
62+
await wrappedHandler(request, env, ctx)
63+
throw new Error('Expected error was not thrown')
64+
} catch (err) {
65+
expect(rateLimiter.limit).toHaveBeenCalledTimes(1)
66+
expect(rateLimiter.limit).toHaveBeenCalledWith({ key: ctx.dataCid.toString() })
67+
expect(handler).not.toHaveBeenCalled()
68+
expect(err).toBeInstanceOf(HttpError)
69+
expect(err.message).toBe('Too Many Requests')
70+
}
71+
})
72+
73+
it('should call next if auth token is present but no token metadata and rate limit is not exceeded', async () => {
2274
rateLimiter.limit.mockResolvedValue({ success: true })
75+
env.AUTH_TOKEN_METADATA.get.mockResolvedValue(null)
2376

2477
const request = {
2578
headers: {
26-
get: jest.fn().mockReturnValue(null) // No Authorization header
79+
get: jest.fn((header) => {
80+
if (header === 'Authorization') {
81+
return 'test-token'
82+
}
83+
return null
84+
})
2785
}
2886
}
29-
const cid = 'test-cid' // TODO: figure out how to generate a CID for testing
30-
const ctx = { dataCid: cid }
87+
const ctx = { dataCid: 'test-cid' }
3188
const wrappedHandler = withRateLimits(handler)
3289

3390
await wrappedHandler(request, env, ctx)
3491

3592
expect(rateLimiter.limit).toHaveBeenCalledTimes(1)
36-
expect(rateLimiter.limit).toHaveBeenCalledWith({ key: cid })
93+
expect(rateLimiter.limit).toHaveBeenCalledWith({ key: ctx.dataCid.toString() })
3794
expect(handler).toHaveBeenCalledTimes(1)
3895
expect(handler).toHaveBeenCalledWith(request, env, ctx)
3996
})
4097

41-
it('should throw an error if rate limit is exceeded', async () => {
98+
it('should throw an error if auth token is present but no token metadata and rate limit is exceeded', async () => {
4299
rateLimiter.limit.mockResolvedValue({ success: false })
100+
env.AUTH_TOKEN_METADATA.get.mockResolvedValue(null)
43101

44102
const request = {
45103
headers: {
46-
get: jest.fn().mockReturnValue(null) // No Authorization header
104+
get: jest.fn((header) => {
105+
if (header === 'Authorization') {
106+
return 'test-token'
107+
}
108+
return null
109+
})
47110
}
48111
}
49-
const cid = 'test-cid' // TODO: figure out how to generate a CID for testing
50-
const ctx = { dataCid: cid }
112+
const ctx = { dataCid: 'test-cid' }
51113
const wrappedHandler = withRateLimits(handler)
52114

53115
try {
54116
await wrappedHandler(request, env, ctx)
55117
throw new Error('Expected error was not thrown')
56118
} catch (err) {
57119
expect(rateLimiter.limit).toHaveBeenCalledTimes(1)
58-
expect(rateLimiter.limit).toHaveBeenCalledWith({ key: cid })
120+
expect(rateLimiter.limit).toHaveBeenCalledWith({ key: ctx.dataCid.toString() })
59121
expect(handler).not.toHaveBeenCalled()
60122
expect(err).toBeInstanceOf(HttpError)
61123
expect(err.message).toBe('Too Many Requests')
62124
}
63125
})
126+
127+
it('should call next if auth token is present and token metadata is invalid but rate limit is not exceeded', async () => {
128+
rateLimiter.limit.mockResolvedValue({ success: true })
129+
env.AUTH_TOKEN_METADATA.get.mockResolvedValue(JSON.stringify({ invalid: true }))
130+
131+
const request = {
132+
headers: {
133+
get: jest.fn((header) => {
134+
if (header === 'Authorization') {
135+
return 'test-token'
136+
}
137+
return null
138+
})
139+
}
140+
}
141+
const ctx = { dataCid: 'test-cid' }
142+
const wrappedHandler = withRateLimits(handler)
143+
144+
await wrappedHandler(request, env, ctx)
145+
146+
expect(rateLimiter.limit).toHaveBeenCalledTimes(1)
147+
expect(rateLimiter.limit).toHaveBeenCalledWith({ key: ctx.dataCid.toString() })
148+
expect(handler).toHaveBeenCalledTimes(1)
149+
expect(handler).toHaveBeenCalledWith(request, env, ctx)
150+
})
151+
152+
it('should throw an error if auth token is present and token metadata is invalid and rate limit is exceeded', async () => {
153+
rateLimiter.limit.mockResolvedValue({ success: false })
154+
env.AUTH_TOKEN_METADATA.get.mockResolvedValue(JSON.stringify({ invalid: true }))
155+
156+
const request = {
157+
headers: {
158+
get: jest.fn((header) => {
159+
if (header === 'Authorization') {
160+
return 'test-token'
161+
}
162+
return null
163+
})
164+
}
165+
}
166+
const ctx = { dataCid: 'test-cid' }
167+
const wrappedHandler = withRateLimits(handler)
168+
169+
try {
170+
await wrappedHandler(request, env, ctx)
171+
throw new Error('Expected error was not thrown')
172+
} catch (err) {
173+
expect(rateLimiter.limit).toHaveBeenCalledTimes(1)
174+
expect(rateLimiter.limit).toHaveBeenCalledWith({ key: ctx.dataCid.toString() })
175+
expect(handler).not.toHaveBeenCalled()
176+
expect(err).toBeInstanceOf(HttpError)
177+
expect(err.message).toBe('Too Many Requests')
178+
}
179+
})
180+
181+
it('should call next if auth token is present and token metadata is valid', async () => {
182+
env.AUTH_TOKEN_METADATA.get.mockResolvedValue(JSON.stringify({ invalid: false }))
183+
184+
const request = {
185+
headers: {
186+
get: jest.fn((header) => {
187+
if (header === 'Authorization') {
188+
return 'test-token'
189+
}
190+
return null
191+
})
192+
}
193+
}
194+
const ctx = { dataCid: 'test-cid' }
195+
const wrappedHandler = withRateLimits(handler)
196+
197+
await wrappedHandler(request, env, ctx)
198+
199+
expect(handler).toHaveBeenCalledTimes(1)
200+
expect(handler).toHaveBeenCalledWith(request, env, ctx)
201+
})
64202
})

0 commit comments

Comments
 (0)