calendar/packages/features/ee/api-keys/lib/autoLock.test.ts

import type { RatelimitResponse } from "@unkey/ratelimit";
import { vi, describe, it, expect, beforeEach, afterEach } from "vitest";

import { hashAPIKey } from "@calcom/features/ee/api-keys/lib/apiKeys";
import { RedisService } from "@calcom/features/redis/RedisService";
import prisma from "@calcom/prisma";

import { handleAutoLock } from "./autoLock";

// Mock the dependencies
vi.mock("@calcom/features/redis/RedisService");
vi.mock("@calcom/features/ee/api-keys/lib/apiKeys", () => ({
  hashAPIKey: vi.fn((key) => `hashed_${key}`),
}));
vi.mock("@calcom/prisma", () => ({
  default: {
    user: {
      update: vi.fn(),
    },
    apiKey: {
      findUnique: vi.fn(),
    },
  },
}));

describe("autoLock", () => {
  const mockRedis = {
    get: vi.fn(),
    set: vi.fn(),
    del: vi.fn(),
    expire: vi.fn(),
  };

  beforeEach(() => {
    // Reset all mocks before each test
    vi.clearAllMocks();
    vi.mocked(RedisService).mockImplementation(function () {
      return mockRedis as any;
    });

    // Mock environment variables
    process.env.UPSTASH_REDIS_REST_TOKEN = "test-token";
    process.env.UPSTASH_REDIS_REST_URL = "test-url";
  });

  afterEach(() => {
    delete process.env.UPSTASH_REDIS_REST_TOKEN;
    delete process.env.UPSTASH_REDIS_REST_URL;
  });

  describe("handleAutoLock", () => {
    it("should return early if Upstash env variables are not set", async () => {
      delete process.env.UPSTASH_REDIS_REST_TOKEN;

      const rateLimitResponse: RatelimitResponse = {
        success: false,
        remaining: 0,
        limit: 5,
        reset: 0,
      };

      await handleAutoLock({
        identifier: "test@example.com",
        identifierType: "email",
        rateLimitResponse,
      });
      expect(mockRedis.get).not.toHaveBeenCalled();
    });

    it("should handle identifier with keyword correctly", async () => {
      const rateLimitResponse: RatelimitResponse = {
        success: false,
        remaining: 0,
        limit: 5,
        reset: 0,
      };

      mockRedis.get.mockResolvedValue("0");

      await handleAutoLock({
        identifier: "addSecondaryEmail.test@example.com",
        identifierType: "email",
        rateLimitResponse,
        identifierKeyword: "addSecondaryEmail",
      });

      expect(mockRedis.get).toHaveBeenCalledWith("autolock:email:addSecondaryEmail:test@example.com.count");
    });

    it("should increment counter when below threshold", async () => {
      const rateLimitResponse: RatelimitResponse = {
        success: false,
        remaining: 0,
        limit: 5,
        reset: 0,
      };

      mockRedis.get.mockResolvedValue("2");

      await handleAutoLock({
        identifier: "test@example.com",
        identifierType: "email",
        rateLimitResponse,
      });

      expect(mockRedis.set).toHaveBeenCalledWith("autolock:email:test@example.com.count", "3");
      expect(mockRedis.expire).toHaveBeenCalledWith("autolock:email:test@example.com.count", 1800);
      expect(prisma.user.update).not.toHaveBeenCalled();
    });

    it("should lock user when threshold is reached", async () => {
      const rateLimitResponse: RatelimitResponse = {
        success: false,
        remaining: 0,
        limit: 5,
        reset: 0,
      };

      mockRedis.get.mockResolvedValue("4");

      await handleAutoLock({
        identifier: "test@example.com",
        identifierType: "email",
        rateLimitResponse,
      });

      expect(prisma.user.update).toHaveBeenCalledWith({
        where: { email: "test@example.com" },
        data: { locked: true },
        select: {
          id: true,
          email: true,
          username: true,
        },
      });
      expect(mockRedis.del).toHaveBeenCalledWith("autolock:email:test@example.com.count");
    });

    it("should respect custom threshold and duration", async () => {
      const rateLimitResponse: RatelimitResponse = {
        success: false,
        remaining: 0,
        limit: 5,
        reset: 0,
      };

      mockRedis.get.mockResolvedValue("1");

      await handleAutoLock({
        identifier: "test@example.com",
        identifierType: "email",
        rateLimitResponse,
        autolockThreshold: 3,
        autolockDuration: 30 * 60 * 1000, // 30 minutes
      });

      expect(mockRedis.set).toHaveBeenCalledWith("autolock:email:test@example.com.count", "2");
      expect(mockRedis.expire).toHaveBeenCalledWith("autolock:email:test@example.com.count", 1800);
      expect(prisma.user.update).not.toHaveBeenCalled();
    });

    it("should lock API key and associated user", async () => {
      const rateLimitResponse: RatelimitResponse = {
        success: false,
        remaining: 0,
        limit: 5,
        reset: 0,
      };

      mockRedis.get.mockResolvedValue("4");
      const testApiKey = "test_api_key_123";

      // Mock API key lookup
      vi.mocked(prisma.apiKey.findUnique).mockResolvedValue({
        hashedKey: `hashed_${testApiKey}`,
        user: { id: 456 },
      } as any);

      await handleAutoLock({
        identifier: testApiKey,
        identifierType: "apiKey",
        rateLimitResponse,
      });

      // Verify the API key was hashed
      expect(hashAPIKey).toHaveBeenCalledWith(testApiKey);

      // Verify the associated user was locked
      expect(prisma.user.update).toHaveBeenCalledWith({
        where: { id: 456 },
        data: { locked: true },
        select: {
          id: true,
          email: true,
          username: true,
        },
      });
    });

    it("should throw error when API key has no associated user", async () => {
      const rateLimitResponse: RatelimitResponse = {
        success: false,
        remaining: 0,
        limit: 5,
        reset: 0,
      };

      mockRedis.get.mockResolvedValue("4"); // Over threshold to trigger lock
      const testApiKey = "test_api_key_123";

      // Mock API key lookup with no user
      vi.mocked(prisma.apiKey.findUnique).mockResolvedValue(null);

      // Expect handleAutoLock to throw the error from lockUser
      await expect(async () => {
        await handleAutoLock({
          identifier: testApiKey,
          identifierType: "apiKey",
          rateLimitResponse,
        });
      }).rejects.toThrow("No user found for this API key.");
    });

    it("should not increment counter when rate limit is successful", async () => {
      const rateLimitResponse: RatelimitResponse = {
        success: true,
        remaining: 5,
        limit: 5,
        reset: 0,
      };

      await handleAutoLock({
        identifier: "test@example.com",
        identifierType: "email",
        rateLimitResponse,
      });

      expect(mockRedis.get).not.toHaveBeenCalled();
      expect(mockRedis.set).not.toHaveBeenCalled();
    });

    it("should initialize counter when it doesn't exist", async () => {
      const rateLimitResponse: RatelimitResponse = {
        success: false,
        remaining: 0,
        limit: 5,
        reset: 0,
      };

      mockRedis.get.mockResolvedValue(null);

      await handleAutoLock({
        identifier: "test@example.com",
        identifierType: "email",
        rateLimitResponse,
      });

      expect(mockRedis.set).toHaveBeenCalledWith("autolock:email:test@example.com.count", "1");
      expect(mockRedis.expire).toHaveBeenCalledWith("autolock:email:test@example.com.count", 1800);
    });

    it("should handle Redis errors gracefully", async () => {
      const rateLimitResponse: RatelimitResponse = {
        success: false,
        remaining: 0,
        limit: 5,
        reset: 0,
      };

      mockRedis.get.mockRejectedValue(new Error("Redis connection error"));

      const result = await handleAutoLock({
        identifier: "test@example.com",
        identifierType: "email",
        rateLimitResponse,
      });

      expect(result).toBe(false);
    });
  });

  describe("lockUser", () => {
    it("should lock user by email", async () => {
      const rateLimitResponse: RatelimitResponse = {
        success: false,
        remaining: 0,
        limit: 5,
        reset: 0,
      };

      mockRedis.get.mockResolvedValue("4");

      await handleAutoLock({
        identifier: "test@example.com",
        identifierType: "email",
        rateLimitResponse,
      });

      expect(prisma.user.update).toHaveBeenCalledWith({
        where: { email: "test@example.com" },
        data: { locked: true },
        select: {
          id: true,
          email: true,
          username: true,
        },
      });
    });

    it("should lock user by userId", async () => {
      const rateLimitResponse: RatelimitResponse = {
        success: false,
        remaining: 0,
        limit: 5,
        reset: 0,
      };

      mockRedis.get.mockResolvedValue("4");

      await handleAutoLock({
        identifier: "123",
        identifierType: "userId",
        rateLimitResponse,
      });

      expect(prisma.user.update).toHaveBeenCalledWith({
        where: { id: 123 },
        data: { locked: true },
        select: {
          id: true,
          email: true,
          username: true,
        },
      });
    });
  });
});