testing is done, fixing algo now

Author: Yufei Wu

src/rateLimiters/fixedWindow.ts

@@ -70,20 +70,21 @@ class FixedWindow implements RateLimiter {
                 await this.client.setex(uuid, keyExpiry, JSON.stringify(newUserWindow));
                 return { success: true, tokens: this.capacity - newUserWindow.currentTokens };
             }
-
+            // optional
             await this.client.setex(uuid, keyExpiry, JSON.stringify(newUserWindow));
+            return { success: false, tokens: this.capacity };
         }
 
         const window: RedisWindow = await JSON.parse(windowJSON as string);
 
         const updatedUserWindow = this.updateTimeWindow(window, timestamp);
-        if (window.currentTokens > this.capacity) {
+        if (tokens > this.capacity) {
             await this.client.setex(uuid, keyExpiry, JSON.stringify(updatedUserWindow));
-            return { success: false, tokens: window.currentTokens };
+            return { success: false, tokens: this.capacity - updatedUserWindow.currentTokens };
         }
         await this.client.setex(uuid, keyExpiry, JSON.stringify(updatedUserWindow));
 
-        return { success: false, tokens: updatedUserWindow.currentTokens };
+        return { success: true, tokens: updatedUserWindow.currentTokens };
     }
 
     /**
@@ -98,7 +99,8 @@ class FixedWindow implements RateLimiter {
             currentTokens: window.currentTokens,
             fixedWindowStart: window.fixedWindowStart,
         };
-        if (timestamp > window.fixedWindowStart + this.windowSize) {
+        if (timestamp >= window.fixedWindowStart + this.windowSize) {
+            // TODO: should base on user's current timestamp
             updatedUserWindow.fixedWindowStart = window.fixedWindowStart + this.windowSize;
             updatedUserWindow.currentTokens = 0;
         }

test/rateLimiters/fixedWindow.test.ts

@@ -26,10 +26,10 @@ async function getWindowFromClient(redisClient: ioredis.Redis, uuid: string): Pr
 async function setTokenCountInClient(
     redisClient: ioredis.Redis,
     uuid: string,
-    tokens: number,
-    time: number
+    currentTokens: number,
+    fixedWindowStart: number
 ) {
-    const value: RedisWindow = { currentTokens: tokens, fixedWindowStart: time };
+    const value: RedisWindow = { currentTokens, fixedWindowStart };
     await redisClient.set(uuid, JSON.stringify(value));
 }
 describe('Test FixedWindow Rate Limiter', () => {
@@ -50,7 +50,84 @@ describe('Test FixedWindow Rate Limiter', () => {
                     CAPACITY - withdraw5
                 );
                 const tokenCountFull = await getWindowFromClient(client, user1);
-                expect(tokenCountFull.currentTokens).toBe(CAPACITY - withdraw5);
+                expect(tokenCountFull.currentTokens).toBe(5);
+            });
+            test('reached 40% capacity in current time window and still can pass request', async () => {
+                const initial = 5;
+                const partialWithdraw = 2;
+                expect(
+                    (await limiter.processRequest(user2, timestamp, initial + partialWithdraw))
+                        .tokens
+                ).toBe(CAPACITY - initial - partialWithdraw);
+                // TODO: using setWindowClient first and update the capacity (not working)
+                // TODO: might need to look into why is not taking second request...
+
+                const tokenCountPartial = await getWindowFromClient(client, user2);
+                expect(tokenCountPartial.currentTokens).toBe(initial + partialWithdraw);
+            });
+
+            // TODO: need to fix processRequest Fn()
+
+            xtest('window is partially full and request has no leftover tokens', async () => {
+                const initial = 6;
+                const partialWithdraw = 5;
+                await setTokenCountInClient(client, user2, initial, timestamp);
+                expect(
+                    (await limiter.processRequest(user2, timestamp, partialWithdraw)).success
+                ).toBe(false);
+                const tokenCountPartialToEmpty = await getWindowFromClient(client, user2);
+                expect(tokenCountPartialToEmpty.currentTokens).toBe(10);
+            });
+        });
+        describe('after a BLOCKED request...', () => {
+            afterEach(() => {
+                client.flushall();
+            });
+            test('initial request is greater than capacity', async () => {
+                // expect remaining tokens to be 10, b/c the 11 token request should be blocked
+                expect((await limiter.processRequest(user1, timestamp, 11)).success).toBe(false);
+                // expect current tokens in the window to still be 0
+                expect((await getWindowFromClient(client, user1)).currentTokens).toBe(0);
+            });
+            xtest('window is partially full but not enough time elapsed to reach new window', async () => {
+                const requestedTokens = 9;
+
+                await setTokenCountInClient(client, user2, requestedTokens, timestamp);
+                // expect remaining tokens to be 1, b/c the 2-token-request should be blocked
+                // TODO: fix processRequest function
+                const result = await limiter.processRequest(user2, timestamp + WINDOW_SIZE - 1, 2);
+
+                expect(result.success).toBe(false);
+                // TODO: fix this, (expect 1, but get 9 -- which is # of currentTokens)
+                expect(result.tokens).toBe(CAPACITY - requestedTokens);
+
+                // expect current tokens in the window to still be 9
+                expect((await getWindowFromClient(client, user2)).currentTokens).toBe(9);
+            });
+        });
+        describe('updateTimeWindow function works as expect', () => {
+            afterEach(() => {
+                client.flushall();
+            });
+            xtest('New window is initialized after reaching the window size', async () => {
+                const fullRequest = 10;
+                await setTokenCountInClient(client, user3, fullRequest, timestamp);
+                const noAccess = await limiter.processRequest(
+                    user3,
+                    timestamp + WINDOW_SIZE - 1,
+                    1
+                );
+
+                // TODO: fix processRequest function
+                // expect cannot pass any request
+                expect(noAccess.success).toBe(false);
+
+                const newRequest = 3;
+                await setTokenCountInClient(client, user3, newRequest, timestamp + WINDOW_SIZE + 1);
+                const newAccess = await limiter.processRequest(user3, timestamp, 1);
+
+                expect(newAccess.success).toBe(true);
+                expect(newAccess.tokens).toBe(4);
             });
         });
     });