bypass rate limiter in unit tests

Author: devksingh4

src/api/functions/rateLimit.ts

@@ -0,0 +1,66 @@
+import {
+  ConditionalCheckFailedException,
+  UpdateItemCommand,
+} from "@aws-sdk/client-dynamodb";
+import { DynamoDBClient } from "@aws-sdk/client-dynamodb";
+import { genericConfig } from "common/config.js";
+
+interface RateLimitParams {
+  ddbClient: DynamoDBClient;
+  rateLimitIdentifier: string;
+  duration: number;
+  limit: number;
+  userIdentifier: string;
+}
+
+export async function isAtLimit({
+  ddbClient,
+  rateLimitIdentifier,
+  duration,
+  limit,
+  userIdentifier,
+}: RateLimitParams): Promise<{
+  limited: boolean;
+  resetTime: number;
+  used: number;
+}> {
+  const nowInSeconds = Math.floor(Date.now() / 1000);
+  const timeWindow = Math.floor(nowInSeconds / duration) * duration;
+  const PK = `rate-limit:${rateLimitIdentifier}:${userIdentifier}:${timeWindow}`;
+
+  try {
+    const result = await ddbClient.send(
+      new UpdateItemCommand({
+        TableName: genericConfig.RateLimiterDynamoTableName,
+        Key: {
+          PK: { S: PK },
+          SK: { S: "counter" },
+        },
+        UpdateExpression: "ADD #rateLimitCount :inc SET #ttl = :ttl",
+        ConditionExpression:
+          "attribute_not_exists(#rateLimitCount) OR #rateLimitCount <= :limit",
+        ExpressionAttributeValues: {
+          ":inc": { N: "1" },
+          ":limit": { N: limit.toString() },
+          ":ttl": { N: (timeWindow + duration).toString() },
+        },
+        ExpressionAttributeNames: {
+          "#rateLimitCount": "rateLimitCount",
+          "#ttl": "ttl",
+        },
+        ReturnValues: "UPDATED_NEW",
+        ReturnValuesOnConditionCheckFailure: "ALL_OLD",
+      }),
+    );
+    return {
+      limited: false,
+      used: parseInt(result.Attributes?.rateLimitCount.N || "1", 10),
+      resetTime: timeWindow + duration,
+    };
+  } catch (error) {
+    if (error instanceof ConditionalCheckFailedException) {
+      return { limited: true, resetTime: timeWindow + duration, used: limit };
+    }
+    throw error;
+  }
+}

src/api/plugins/rateLimiter.ts

@@ -1,10 +1,5 @@
 import fp from "fastify-plugin";
-import {
-  ConditionalCheckFailedException,
-  UpdateItemCommand,
-} from "@aws-sdk/client-dynamodb";
-import { DynamoDBClient } from "@aws-sdk/client-dynamodb";
-import { genericConfig } from "common/config.js";
+import { isAtLimit } from "api/functions/rateLimit.js";
 import { FastifyPluginAsync, FastifyRequest, FastifyReply } from "fastify";
 
 interface RateLimiterOptions {
@@ -13,66 +8,6 @@ interface RateLimiterOptions {
   rateLimitIdentifier?: string;
 }
 
-interface RateLimitParams {
-  ddbClient: DynamoDBClient;
-  rateLimitIdentifier: string;
-  duration: number;
-  limit: number;
-  userIdentifier: string;
-}
-
-async function isAtLimit({
-  ddbClient,
-  rateLimitIdentifier,
-  duration,
-  limit,
-  userIdentifier,
-}: RateLimitParams): Promise<{
-  limited: boolean;
-  resetTime: number;
-  used: number;
-}> {
-  const nowInSeconds = Math.floor(Date.now() / 1000);
-  const timeWindow = Math.floor(nowInSeconds / duration) * duration;
-  const PK = `rate-limit:${rateLimitIdentifier}:${userIdentifier}:${timeWindow}`;
-
-  try {
-    const result = await ddbClient.send(
-      new UpdateItemCommand({
-        TableName: genericConfig.RateLimiterDynamoTableName,
-        Key: {
-          PK: { S: PK },
-          SK: { S: "counter" },
-        },
-        UpdateExpression: "ADD #rateLimitCount :inc SET #ttl = :ttl",
-        ConditionExpression:
-          "attribute_not_exists(#rateLimitCount) OR #rateLimitCount <= :limit",
-        ExpressionAttributeValues: {
-          ":inc": { N: "1" },
-          ":limit": { N: limit.toString() },
-          ":ttl": { N: (timeWindow + duration).toString() },
-        },
-        ExpressionAttributeNames: {
-          "#rateLimitCount": "rateLimitCount",
-          "#ttl": "ttl",
-        },
-        ReturnValues: "UPDATED_NEW",
-        ReturnValuesOnConditionCheckFailure: "ALL_OLD",
-      }),
-    );
-    return {
-      limited: false,
-      used: parseInt(result.Attributes?.rateLimitCount.N || "1", 10),
-      resetTime: timeWindow + duration,
-    };
-  } catch (error) {
-    if (error instanceof ConditionalCheckFailedException) {
-      return { limited: true, resetTime: timeWindow + duration, used: limit };
-    }
-    throw error;
-  }
-}
-
 const rateLimiterPlugin: FastifyPluginAsync<RateLimiterOptions> = async (
   fastify,
   options,

tests/unit/vitest.setup.ts

@@ -2,13 +2,13 @@ import { vi } from "vitest";
 import { allAppRoles, AppRoles } from "../../src/common/roles.js";
 
 vi.mock(
-  import("../../src/api/plugins/rateLimiter.js"),
+  import("../../src/api/functions/rateLimit.js"),
   async (importOriginal) => {
     const mod = await importOriginal();
     return {
       ...mod,
       isAtLimit: vi.fn(async (_) => {
-        return false;
+        return { limited: false, resetTime: 0, used: 1 };
       }),
     };
   },