Nick: expire tests

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();
   });
 });

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<number> {
   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;