Skip to content

Commit cd0f1c1

Browse files
shalarewiczshalarewicz
committed
separated allowed/blocked token bucket tests for clarity
1 parent e88e0c1 commit cd0f1c1

File tree

1 file changed

+73
-57
lines changed

1 file changed

+73
-57
lines changed

‎test/rateLimiters/tokenBucket.test.ts‎

Lines changed: 73 additions & 57 deletions
Original file line numberDiff line numberDiff line change
@@ -42,69 +42,85 @@ xdescribe('Test TokenBucket Rate Limiter', () => {
4242
});
4343

4444
describe('TokenBucket returns correct number of tokens and updates redis store as expected', () => {
45-
test('after an ALLOWED request', async () => {
46-
// Bucket intially full
47-
const withdraw5 = 5;
48-
expect((await limiter.processRequest(user1, timestamp, withdraw5)).tokens).toBe(
49-
CAPACITY - withdraw5
50-
);
51-
const tokenCountFull = await getBucketFromClient(client, user1);
52-
expect(tokenCountFull).toBe(CAPACITY - withdraw5);
53-
54-
// Bucket partially full but enough time has elapsed to fill the bucket since the last request and
55-
// has leftover tokens after reqeust
56-
const initial = 6;
57-
const partialWithdraw = 1;
58-
await setTokenCountInClient(client, user2, initial, timestamp);
59-
expect(
60-
(
61-
await limiter.processRequest(
62-
user2,
63-
timestamp + 1000 * (CAPACITY - initial),
64-
initial + partialWithdraw
65-
)
66-
).tokens
67-
).toBe(CAPACITY - (initial + partialWithdraw));
68-
const tokenCountPartial = await getBucketFromClient(client, user2);
69-
expect(tokenCountPartial).toBe(CAPACITY - (initial + partialWithdraw));
45+
describe('after an ALLOWED request...', () => {
46+
test('bucket is initially full', async () => {
47+
// Bucket intially full
48+
const withdraw5 = 5;
49+
expect((await limiter.processRequest(user1, timestamp, withdraw5)).tokens).toBe(
50+
CAPACITY - withdraw5
51+
);
52+
const tokenCountFull = await getBucketFromClient(client, user1);
53+
expect(tokenCountFull).toBe(CAPACITY - withdraw5);
54+
});
55+
56+
test('bucket is partially full and request has leftover tokens', async () => {
57+
// Bucket partially full but enough time has elapsed to fill the bucket since the last request and
58+
// has leftover tokens after reqeust
59+
const initial = 6;
60+
const partialWithdraw = 1;
61+
await setTokenCountInClient(client, user2, initial, timestamp);
62+
expect(
63+
(
64+
await limiter.processRequest(
65+
user2,
66+
timestamp + 1000 * (CAPACITY - initial),
67+
initial + partialWithdraw
68+
)
69+
).tokens
70+
).toBe(CAPACITY - (initial + partialWithdraw));
71+
const tokenCountPartial = await getBucketFromClient(client, user2);
72+
expect(tokenCountPartial).toBe(CAPACITY - (initial + partialWithdraw));
73+
});
7074

7175
// Bucket partially full and no leftover tokens after reqeust
72-
const initial2 = 6;
73-
await setTokenCountInClient(client, user2, initial, timestamp);
74-
expect((await limiter.processRequest(user2, timestamp, initial2)).tokens).toBe(0);
75-
const tokenCountPartialToEmpty = await getBucketFromClient(client, user2);
76-
expect(tokenCountPartialToEmpty).toBe(0);
77-
78-
// Bucket initially empty but enough time elapsed to paritally fill bucket since last request
79-
await setTokenCountInClient(client, user4, 0, timestamp);
80-
expect((await limiter.processRequest(user4, timestamp + 6000, 4)).tokens).toBe(2);
81-
const count = await getBucketFromClient(client, user4);
82-
expect(count).toBe(2);
76+
test('bucket is partially full and request has no leftover tokens', async () => {
77+
const initial = 6;
78+
await setTokenCountInClient(client, user2, initial, timestamp);
79+
expect((await limiter.processRequest(user2, timestamp, initial)).tokens).toBe(0);
80+
const tokenCountPartialToEmpty = await getBucketFromClient(client, user2);
81+
expect(tokenCountPartialToEmpty).toBe(0);
82+
83+
// Bucket initially empty but enough time elapsed to paritally fill bucket since last request
84+
await setTokenCountInClient(client, user4, 0, timestamp);
85+
expect((await limiter.processRequest(user4, timestamp + 6000, 4)).tokens).toBe(2);
86+
const count = await getBucketFromClient(client, user4);
87+
expect(count).toBe(2);
88+
});
8389
});
8490

85-
test('after a BLOCKED request', async () => {
91+
describe('after a BLOCKED request...', () => {
8692
let redisData: RedisBucket;
87-
// Initial request greater than capacity
88-
expect((await limiter.processRequest(user1, timestamp, CAPACITY + 1)).tokens).toBe(
89-
CAPACITY
90-
);
9193

92-
redisData = await getBucketFromClient(client, user1);
93-
expect(redisData.tokens).toBe(CAPACITY);
94-
95-
// Bucket is partially full and time has elapsed but not enough to allow the current request
96-
const fillLevel = 5;
97-
const timeDelta = 3;
98-
const requestedTokens = 9;
99-
await setTokenCountInClient(client, user2, fillLevel, timestamp);
100-
101-
expect(
102-
(await limiter.processRequest(user1, timestamp + timeDelta * 1000, requestedTokens))
103-
.tokens
104-
).toBe(fillLevel + timeDelta * REFILL_RATE);
105-
106-
redisData = await getBucketFromClient(client, user2);
107-
expect(redisData.tokens).toBe(fillLevel + timeDelta * REFILL_RATE);
94+
test('where intial request is greater than bucket capacity', async () => {
95+
// Initial request greater than capacity
96+
expect((await limiter.processRequest(user1, timestamp, CAPACITY + 1)).tokens).toBe(
97+
CAPACITY
98+
);
99+
100+
redisData = await getBucketFromClient(client, user1);
101+
expect(redisData.tokens).toBe(CAPACITY);
102+
});
103+
104+
test('Bucket is partially full but not enough time elapsed to complete the request', async () => {
105+
// Bucket is partially full and time has elapsed but not enough to allow the current request
106+
const fillLevel = 5;
107+
const timeDelta = 3;
108+
const requestedTokens = 9;
109+
await setTokenCountInClient(client, user2, fillLevel, timestamp);
110+
111+
expect(
112+
(
113+
await limiter.processRequest(
114+
user1,
115+
timestamp + timeDelta * 1000,
116+
requestedTokens
117+
)
118+
).tokens
119+
).toBe(fillLevel + timeDelta * REFILL_RATE);
120+
121+
redisData = await getBucketFromClient(client, user2);
122+
expect(redisData.tokens).toBe(fillLevel + timeDelta * REFILL_RATE);
123+
});
108124
});
109125
});
110126

0 commit comments

Comments
 (0)