ShieldAI/services/darkwatch/test/scheduler.test.ts

import { describe, it, expect, beforeEach, afterEach } from "vitest";
import { ScanScheduler } from "../src/scheduler/ScanScheduler";
import { SubscriptionTier } from "@shieldai/types";
import prisma from "@shieldai/db";

let runId = Date.now();

describe("ScanScheduler", () => {
  let scheduler: ScanScheduler;

  beforeEach(() => {
    scheduler = new ScanScheduler();
  });

  describe("static tier configuration", () => {
    it("returns correct interval for BASIC tier", () => {
      expect(ScanScheduler.getIntervalForTier(SubscriptionTier.BASIC)).toBe(1440);
    });

    it("returns correct interval for PLUS tier", () => {
      expect(ScanScheduler.getIntervalForTier(SubscriptionTier.PLUS)).toBe(360);
    });

    it("returns correct interval for PREMIUM tier", () => {
      expect(ScanScheduler.getIntervalForTier(SubscriptionTier.PREMIUM)).toBe(60);
    });

    it("returns correct cron for BASIC tier", () => {
      expect(ScanScheduler.getCronForTier(SubscriptionTier.BASIC)).toBe("0 0 * * *");
    });

    it("returns correct cron for PLUS tier", () => {
      expect(ScanScheduler.getCronForTier(SubscriptionTier.PLUS)).toBe("0 */6 * * *");
    });

    it("returns correct cron for PREMIUM tier", () => {
      expect(ScanScheduler.getCronForTier(SubscriptionTier.PREMIUM)).toBe("0 * * * *");
    });
  });

  describe("ensureScheduleForUser", () => {
    let userId: string;

    beforeEach(async () => {
      const user = await prisma.user.create({
        data: {
          email: `scheduler-test-${runId}@shieldai.local`,
          subscriptionTier: "BASIC",
        },
      });
      userId = user.id;
      runId = Date.now();
    });

    afterEach(async () => {
      await prisma.scanSchedule.deleteMany({ where: { userId } });
      await prisma.user.delete({ where: { id: userId } });
    });

    it("creates schedule for new user", async () => {
      const result = await scheduler.ensureScheduleForUser(userId);
      expect(result.scheduled).toBe(true);
      expect(result.intervalMinutes).toBe(1440);

      const schedule = await scheduler.getSchedule(userId);
      expect(schedule).not.toBeNull();
      expect(schedule?.status).toBe("ACTIVE");
      expect(schedule?.cronExpression).toBe("0 0 * * *");
    });

    it("updates schedule on tier change", async () => {
      await scheduler.ensureScheduleForUser(userId);

      await prisma.user.update({
        where: { id: userId },
        data: { subscriptionTier: "PREMIUM" },
      });

      const result = await scheduler.ensureScheduleForUser(userId);
      expect(result.intervalMinutes).toBe(60);

      const schedule = await scheduler.getSchedule(userId);
      expect(schedule?.cronExpression).toBe("0 * * * *");
    });

    it("returns false for non-existent user", async () => {
      const result = await scheduler.ensureScheduleForUser("non-existent-id");
      expect(result.scheduled).toBe(false);
      expect(result.intervalMinutes).toBe(0);
    });
  });

  describe("schedule lifecycle", () => {
    let userId: string;

    beforeEach(async () => {
      const user = await prisma.user.create({
        data: {
          email: `lifecycle-test-${runId}@shieldai.local`,
          subscriptionTier: "PLUS",
        },
      });
      userId = user.id;
      runId = Date.now();
      await scheduler.ensureScheduleForUser(userId);
    });

    afterEach(async () => {
      await prisma.scanSchedule.deleteMany({ where: { userId } });
      await prisma.user.delete({ where: { id: userId } });
    });

    it("marks schedule as scanned and updates next scan time", async () => {
      const before = await scheduler.getSchedule(userId);
      const nextScan = await scheduler.markScanned(userId);

      const after = await scheduler.getSchedule(userId);
      expect(after?.lastScanAt).not.toBeNull();
      expect(after?.nextScanAt?.getTime()).toBeGreaterThan(nextScan.getTime() - 5000);
      expect(after?.nextScanAt).not.toEqual(before?.nextScanAt);
    });

    it("pauses schedule", async () => {
      await scheduler.pauseSchedule(userId);

      const schedule = await scheduler.getSchedule(userId);
      expect(schedule?.status).toBe("PAUSED");
    });

    it("resumes paused schedule", async () => {
      await scheduler.pauseSchedule(userId);
      await scheduler.resumeSchedule(userId);

      const schedule = await scheduler.getSchedule(userId);
      expect(schedule?.status).toBe("ACTIVE");
    });
  });

  describe("getDueSchedules", () => {
    let userId1: string;
    let userId2: string;

    beforeEach(async () => {
      const user1 = await prisma.user.create({
        data: {
          email: `due-test-1-${runId}@shieldai.local`,
          subscriptionTier: "PREMIUM",
        },
      });
      userId1 = user1.id;

      const user2 = await prisma.user.create({
        data: {
          email: `due-test-2-${runId}@shieldai.local`,
          subscriptionTier: "BASIC",
        },
      });
      userId2 = user2.id;
      runId = Date.now();

      await scheduler.ensureScheduleForUser(userId1);
      await scheduler.ensureScheduleForUser(userId2);
    });

    afterEach(async () => {
      await prisma.scanSchedule.deleteMany({ where: { userId: userId1 } });
      await prisma.scanSchedule.deleteMany({ where: { userId: userId2 } });
      await prisma.user.delete({ where: { id: userId1 } });
      await prisma.user.delete({ where: { id: userId2 } });
    });

    it("returns schedules that are due", async () => {
      const pastDate = new Date(Date.now() - 60000);
      await prisma.scanSchedule.update({
        where: { userId: userId1 },
        data: { nextScanAt: pastDate },
      });

      const due = await scheduler.getDueSchedules();
      const dueUserIds = due.map((s) => s.userId);
      expect(dueUserIds).toContain(userId1);
    });

    it("includes schedules with null nextScanAt", async () => {
      await prisma.scanSchedule.update({
        where: { userId: userId2 },
        data: { nextScanAt: null },
      });

      const due = await scheduler.getDueSchedules();
      const dueUserIds = due.map((s) => s.userId);
      expect(dueUserIds).toContain(userId2);
    });
  });
});