From 3c3412e893d68c089dd2a0d1d4ed69c8f40ba061 Mon Sep 17 00:00:00 2001 From: Nicolas Date: Tue, 16 Jul 2024 22:45:12 -0400 Subject: [PATCH] Update rate-limiter.test.ts --- apps/api/src/services/rate-limiter.test.ts | 315 +++++++++++++++++++-- 1 file changed, 296 insertions(+), 19 deletions(-) diff --git a/apps/api/src/services/rate-limiter.test.ts b/apps/api/src/services/rate-limiter.test.ts index 47a24a7e..c49c85d9 100644 --- a/apps/api/src/services/rate-limiter.test.ts +++ b/apps/api/src/services/rate-limiter.test.ts @@ -1,48 +1,98 @@ -import { getRateLimiter, serverRateLimiter, testSuiteRateLimiter, redisClient } from "./rate-limiter"; +import { + getRateLimiter, + serverRateLimiter, + testSuiteRateLimiter, + redisRateLimitClient, +} from "./rate-limiter"; import { RateLimiterMode } from "../../src/types"; import { RateLimiterRedis } from "rate-limiter-flexible"; describe("Rate Limiter Service", () => { beforeAll(async () => { - await redisClient.connect(); + try { + await redisRateLimitClient.connect(); + // if (process.env.REDIS_RATE_LIMIT_URL === "redis://localhost:6379") { + // console.log("Erasing all keys"); + // // erase all the keys that start with "test-prefix" + // const keys = await redisRateLimitClient.keys("test-prefix:*"); + // if (keys.length > 0) { + // await redisRateLimitClient.del(...keys); + // } + // } + } catch (error) {} }); afterAll(async () => { - await redisClient.disconnect(); + try { + // if (process.env.REDIS_RATE_LIMIT_URL === "redis://localhost:6379") { + await redisRateLimitClient.disconnect(); + // } + } catch (error) {} }); it("should return the testSuiteRateLimiter for specific tokens", () => { - const limiter = getRateLimiter("crawl" as RateLimiterMode, "a01ccae"); + const limiter = getRateLimiter( + "crawl" as RateLimiterMode, + "test-prefix:a01ccae" + ); expect(limiter).toBe(testSuiteRateLimiter); - const limiter2 = getRateLimiter("scrape" as RateLimiterMode, "6254cf9"); + const limiter2 = getRateLimiter( + "scrape" as RateLimiterMode, + "test-prefix:6254cf9" + ); expect(limiter2).toBe(testSuiteRateLimiter); }); it("should return the serverRateLimiter if mode is not found", () => { - const limiter = getRateLimiter("nonexistent" as RateLimiterMode, "someToken"); + const limiter = getRateLimiter( + "nonexistent" as RateLimiterMode, + "test-prefix:someToken" + ); expect(limiter).toBe(serverRateLimiter); }); it("should return the correct rate limiter based on mode and plan", () => { - const limiter = getRateLimiter("crawl" as RateLimiterMode, "someToken", "free"); + const limiter = getRateLimiter( + "crawl" as RateLimiterMode, + "test-prefix:someToken", + "free" + ); expect(limiter.points).toBe(2); - const limiter2 = getRateLimiter("scrape" as RateLimiterMode, "someToken", "standard"); + const limiter2 = getRateLimiter( + "scrape" as RateLimiterMode, + "test-prefix:someToken", + "standard" + ); expect(limiter2.points).toBe(50); - const limiter3 = getRateLimiter("search" as RateLimiterMode, "someToken", "growth"); + const limiter3 = getRateLimiter( + "search" as RateLimiterMode, + "test-prefix:someToken", + "growth" + ); expect(limiter3.points).toBe(500); - const limiter4 = getRateLimiter("crawlStatus" as RateLimiterMode, "someToken", "growth"); + const limiter4 = getRateLimiter( + "crawlStatus" as RateLimiterMode, + "test-prefix:someToken", + "growth" + ); expect(limiter4.points).toBe(150); }); it("should return the default rate limiter if plan is not provided", () => { - const limiter = getRateLimiter("crawl" as RateLimiterMode, "someToken"); + const limiter = getRateLimiter( + "crawl" as RateLimiterMode, + "test-prefix:someToken" + ); expect(limiter.points).toBe(3); - const limiter2 = getRateLimiter("scrape" as RateLimiterMode, "someToken"); + const limiter2 = getRateLimiter( + "scrape" as RateLimiterMode, + "test-prefix:someToken" + ); expect(limiter2.points).toBe(20); }); @@ -50,7 +100,7 @@ describe("Rate Limiter Service", () => { const keyPrefix = "test-prefix"; const points = 10; const limiter = new RateLimiterRedis({ - storeClient: redisClient, + storeClient: redisRateLimitClient, keyPrefix, points, duration: 60, @@ -62,26 +112,253 @@ describe("Rate Limiter Service", () => { }); it("should return the correct rate limiter for 'preview' mode", () => { - const limiter = getRateLimiter("preview" as RateLimiterMode, "someToken", "free"); + const limiter = getRateLimiter( + "preview" as RateLimiterMode, + "test-prefix:someToken", + "free" + ); expect(limiter.points).toBe(5); - const limiter2 = getRateLimiter("preview" as RateLimiterMode, "someToken"); + const limiter2 = getRateLimiter( + "preview" as RateLimiterMode, + "test-prefix:someToken" + ); expect(limiter2.points).toBe(5); }); it("should return the correct rate limiter for 'account' mode", () => { - const limiter = getRateLimiter("account" as RateLimiterMode, "someToken", "free"); + const limiter = getRateLimiter( + "account" as RateLimiterMode, + "test-prefix:someToken", + "free" + ); expect(limiter.points).toBe(100); - const limiter2 = getRateLimiter("account" as RateLimiterMode, "someToken"); + const limiter2 = getRateLimiter( + "account" as RateLimiterMode, + "test-prefix:someToken" + ); expect(limiter2.points).toBe(100); }); it("should return the correct rate limiter for 'crawlStatus' mode", () => { - const limiter = getRateLimiter("crawlStatus" as RateLimiterMode, "someToken", "free"); + const limiter = getRateLimiter( + "crawlStatus" as RateLimiterMode, + "test-prefix:someToken", + "free" + ); expect(limiter.points).toBe(150); - const limiter2 = getRateLimiter("crawlStatus" as RateLimiterMode, "someToken"); + const limiter2 = getRateLimiter( + "crawlStatus" as RateLimiterMode, + "test-prefix:someToken" + ); expect(limiter2.points).toBe(150); }); + + it("should consume points correctly for 'crawl' mode", async () => { + const limiter = getRateLimiter( + "crawl" as RateLimiterMode, + "test-prefix:someTokenCRAWL", + "free" + ); + const consumePoints = 1; + + const res = await limiter.consume( + "test-prefix:someTokenCRAWL", + consumePoints + ); + expect(res.remainingPoints).toBe(1); + }); + + it("should consume points correctly for 'scrape' mode (DEFAULT)", async () => { + const limiter = getRateLimiter( + "scrape" as RateLimiterMode, + "test-prefix:someTokenX" + ); + const consumePoints = 4; + + const res = await limiter.consume("test-prefix:someTokenX", consumePoints); + expect(res.remainingPoints).toBe(16); + }); + + it("should consume points correctly for 'scrape' mode (HOBBY)", async () => { + const limiter = getRateLimiter( + "scrape" as RateLimiterMode, + "test-prefix:someTokenXY", + "hobby" + ); + // expect hobby to have 100 points + expect(limiter.points).toBe(10); + + const consumePoints = 5; + + const res = await limiter.consume("test-prefix:someTokenXY", consumePoints); + expect(res.consumedPoints).toBe(5); + expect(res.remainingPoints).toBe(5); + }); + + it("should return the correct rate limiter for 'crawl' mode", () => { + const limiter = getRateLimiter( + "crawl" as RateLimiterMode, + "test-prefix:someToken", + "free" + ); + expect(limiter.points).toBe(2); + + const limiter2 = getRateLimiter( + "crawl" as RateLimiterMode, + "test-prefix:someToken", + "starter" + ); + expect(limiter2.points).toBe(3); + + const limiter3 = getRateLimiter( + "crawl" as RateLimiterMode, + "test-prefix:someToken", + "standard" + ); + expect(limiter3.points).toBe(5); + }); + + it("should return the correct rate limiter for 'scrape' mode", () => { + const limiter = getRateLimiter( + "scrape" as RateLimiterMode, + "test-prefix:someToken", + "free" + ); + expect(limiter.points).toBe(5); + + const limiter2 = getRateLimiter( + "scrape" as RateLimiterMode, + "test-prefix:someToken", + "starter" + ); + expect(limiter2.points).toBe(20); + + const limiter3 = getRateLimiter( + "scrape" as RateLimiterMode, + "test-prefix:someToken", + "standard" + ); + expect(limiter3.points).toBe(50); + }); + + it("should return the correct rate limiter for 'search' mode", () => { + const limiter = getRateLimiter( + "search" as RateLimiterMode, + "test-prefix:someToken", + "free" + ); + expect(limiter.points).toBe(5); + + const limiter2 = getRateLimiter( + "search" as RateLimiterMode, + "test-prefix:someToken", + "starter" + ); + expect(limiter2.points).toBe(20); + + const limiter3 = getRateLimiter( + "search" as RateLimiterMode, + "test-prefix:someToken", + "standard" + ); + expect(limiter3.points).toBe(40); + }); + + it("should return the correct rate limiter for 'preview' mode", () => { + const limiter = getRateLimiter( + "preview" as RateLimiterMode, + "test-prefix:someToken", + "free" + ); + expect(limiter.points).toBe(5); + + const limiter2 = getRateLimiter( + "preview" as RateLimiterMode, + "test-prefix:someToken" + ); + expect(limiter2.points).toBe(5); + }); + + it("should return the correct rate limiter for 'account' mode", () => { + const limiter = getRateLimiter( + "account" as RateLimiterMode, + "test-prefix:someToken", + "free" + ); + expect(limiter.points).toBe(100); + + const limiter2 = getRateLimiter( + "account" as RateLimiterMode, + "test-prefix:someToken" + ); + expect(limiter2.points).toBe(100); + }); + + it("should return the correct rate limiter for 'crawlStatus' mode", () => { + const limiter = getRateLimiter( + "crawlStatus" as RateLimiterMode, + "test-prefix:someToken", + "free" + ); + expect(limiter.points).toBe(150); + + const limiter2 = getRateLimiter( + "crawlStatus" as RateLimiterMode, + "test-prefix:someToken" + ); + expect(limiter2.points).toBe(150); + }); + + it("should return the correct rate limiter for 'testSuite' mode", () => { + const limiter = getRateLimiter( + "testSuite" as RateLimiterMode, + "test-prefix:someToken", + "free" + ); + expect(limiter.points).toBe(10000); + + const limiter2 = getRateLimiter( + "testSuite" as RateLimiterMode, + "test-prefix:someToken" + ); + expect(limiter2.points).toBe(10000); + }); + + it("should throw an error when consuming more points than available", async () => { + const limiter = getRateLimiter( + "crawl" as RateLimiterMode, + "test-prefix:someToken" + ); + const consumePoints = limiter.points + 1; + + try { + await limiter.consume("test-prefix:someToken", consumePoints); + } catch (error) { + // expect remaining points to be 0 + const res = await limiter.get("test-prefix:someToken"); + expect(res.remainingPoints).toBe(0); + } + }); + + it("should reset points after duration", async () => { + const keyPrefix = "test-prefix"; + const points = 10; + const duration = 1; // 1 second + const limiter = new RateLimiterRedis({ + storeClient: redisRateLimitClient, + keyPrefix, + points, + duration, + }); + + const consumePoints = 5; + await limiter.consume("test-prefix:someToken", consumePoints); + await new Promise((resolve) => setTimeout(resolve, duration * 1000 + 100)); // Wait for duration + 100ms + + const res = await limiter.consume("test-prefix:someToken", consumePoints); + expect(res.remainingPoints).toBe(points - consumePoints); + }); });