From c009013ff669bd2e6d574fb05a2c36685a9f4630 Mon Sep 17 00:00:00 2001 From: Nicolas Date: Tue, 27 Aug 2024 15:26:43 -0300 Subject: [PATCH] Nick: expire tests --- .../src/lib/__tests__/job-priority.test.ts | 86 ++++++++++++++++--- apps/api/src/lib/job-priority.ts | 18 ++-- 2 files changed, 83 insertions(+), 21 deletions(-) diff --git a/apps/api/src/lib/__tests__/job-priority.test.ts b/apps/api/src/lib/__tests__/job-priority.test.ts index e6f80f05..5a30f92e 100644 --- a/apps/api/src/lib/__tests__/job-priority.test.ts +++ b/apps/api/src/lib/__tests__/job-priority.test.ts @@ -1,4 +1,8 @@ -import { getJobPriority, addJobPriority, deleteJobPriority } from "../job-priority"; +import { + getJobPriority, + addJobPriority, + deleteJobPriority, +} from "../job-priority"; import { redisConnection } from "../../services/queue-service"; import { PlanType } from "../../types"; @@ -7,6 +11,7 @@ jest.mock("../../services/queue-service", () => ({ sadd: jest.fn(), srem: jest.fn(), scard: jest.fn(), + expire: jest.fn(), }, })); @@ -15,18 +20,28 @@ describe("Job Priority Tests", () => { jest.clearAllMocks(); }); - test("addJobPriority should add job_id to the set", async () => { + test("addJobPriority should add job_id to the set and set expiration", async () => { const team_id = "team1"; const job_id = "job1"; await addJobPriority(team_id, job_id); - expect(redisConnection.sadd).toHaveBeenCalledWith(`limit_team_id:${team_id}`, job_id); + expect(redisConnection.sadd).toHaveBeenCalledWith( + `limit_team_id:${team_id}`, + job_id + ); + expect(redisConnection.expire).toHaveBeenCalledWith( + `limit_team_id:${team_id}`, + 60 + ); }); - test("deleteFromSet should remove job_id from the set", async () => { + test("deleteJobPriority should remove job_id from the set", async () => { const team_id = "team1"; const job_id = "job1"; await deleteJobPriority(team_id, job_id); - expect(redisConnection.srem).toHaveBeenCalledWith(`limit_team_id:${team_id}`, job_id); + expect(redisConnection.srem).toHaveBeenCalledWith( + `limit_team_id:${team_id}`, + job_id + ); }); test("getJobPriority should return correct priority based on plan and set length", async () => { @@ -35,11 +50,11 @@ describe("Job Priority Tests", () => { (redisConnection.scard as jest.Mock).mockResolvedValue(150); const priority = await getJobPriority({ plan, team_id }); - expect(priority).toBe(10); + expect(priority).toBe(20); (redisConnection.scard as jest.Mock).mockResolvedValue(250); const priorityExceeded = await getJobPriority({ plan, team_id }); - expect(priorityExceeded).toBe(20); // basePriority + Math.ceil((250 - 200) * 0.2) + expect(priorityExceeded).toBe(40); // basePriority + Math.ceil((250 - 200) * 0.2) }); test("getJobPriority should handle different plans correctly", async () => { @@ -53,9 +68,9 @@ describe("Job Priority Tests", () => { (redisConnection.scard as jest.Mock).mockResolvedValue(150); plan = "hobby"; priority = await getJobPriority({ plan, team_id }); - expect(priority).toBe(35); // basePriority + Math.ceil((150 - 100) * 0.5) + expect(priority).toBe(40); // basePriority + Math.ceil((150 - 50) * 0.3) - (redisConnection.scard as jest.Mock).mockResolvedValue(50); + (redisConnection.scard as jest.Mock).mockResolvedValue(25); plan = "free"; priority = await getJobPriority({ plan, team_id }); expect(priority).toBe(10); @@ -63,6 +78,57 @@ describe("Job Priority Tests", () => { (redisConnection.scard as jest.Mock).mockResolvedValue(60); plan = "free"; priority = await getJobPriority({ plan, team_id }); - expect(priority).toBe(20); // basePriority + Math.ceil((60 - 50) * 1) + expect(priority).toBe(28); // basePriority + Math.ceil((60 - 25) * 0.5) + }); + + test("addJobPriority should reset expiration time when adding new job", async () => { + const team_id = "team1"; + const job_id1 = "job1"; + const job_id2 = "job2"; + + await addJobPriority(team_id, job_id1); + expect(redisConnection.expire).toHaveBeenCalledWith( + `limit_team_id:${team_id}`, + 60 + ); + + // Clear the mock calls + (redisConnection.expire as jest.Mock).mockClear(); + + // Add another job + await addJobPriority(team_id, job_id2); + expect(redisConnection.expire).toHaveBeenCalledWith( + `limit_team_id:${team_id}`, + 60 + ); + }); + + test("Set should expire after 60 seconds", async () => { + const team_id = "team1"; + const job_id = "job1"; + + jest.useFakeTimers(); + + await addJobPriority(team_id, job_id); + expect(redisConnection.expire).toHaveBeenCalledWith( + `limit_team_id:${team_id}`, + 60 + ); + + // Fast-forward time by 59 seconds + jest.advanceTimersByTime(59000); + + // The set should still exist + expect(redisConnection.scard).not.toHaveBeenCalled(); + + // Fast-forward time by 2 more seconds (total 61 seconds) + jest.advanceTimersByTime(2000); + + // Check if the set has been removed (scard should return 0) + (redisConnection.scard as jest.Mock).mockResolvedValue(0); + const setSize = await redisConnection.scard(`limit_team_id:${team_id}`); + expect(setSize).toBe(0); + + jest.useRealTimers(); }); }); diff --git a/apps/api/src/lib/job-priority.ts b/apps/api/src/lib/job-priority.ts index d108938e..361c608d 100644 --- a/apps/api/src/lib/job-priority.ts +++ b/apps/api/src/lib/job-priority.ts @@ -2,19 +2,16 @@ import { redisConnection } from "../../src/services/queue-service"; import { PlanType } from "../../src/types"; import { Logger } from "./logger"; - - -const SET_KEY_PREFIX = "limit_team_id:" +const SET_KEY_PREFIX = "limit_team_id:"; export async function addJobPriority(team_id, job_id) { try { const setKey = SET_KEY_PREFIX + team_id; // Add scrape job id to the set await redisConnection.sadd(setKey, job_id); - + // This approach will reset the expiration time to 60 seconds every time a new job is added to the set. await redisConnection.expire(setKey, 60); - } catch (e) { Logger.error(`Add job priority (sadd) failed: ${team_id}, ${job_id}`); } @@ -34,11 +31,11 @@ export async function deleteJobPriority(team_id, job_id) { export async function getJobPriority({ plan, team_id, - basePriority = 10 + basePriority = 10, }: { plan: PlanType; team_id: string; - basePriority: number; + basePriority?: number; }): Promise { const setKey = SET_KEY_PREFIX + team_id; @@ -52,11 +49,11 @@ export async function getJobPriority({ switch (plan) { case "free": bucketLimit = 25; - planModifier = 1; + planModifier = 0.5; break; case "hobby": bucketLimit = 50; - planModifier = 0.5; + planModifier = 0.3; break; case "standard": case "standardnew": @@ -68,8 +65,7 @@ export async function getJobPriority({ bucketLimit = 200; planModifier = 0.2; break; - - + default: bucketLimit = 25; planModifier = 1;