stack/apps/api/src/prisma/rls-context.provider.spec.ts

import { describe, it, expect, beforeEach, vi } from "vitest";
import { getRlsClient, runWithRlsClient, type TransactionClient } from "./rls-context.provider";

describe("RlsContextProvider", () => {
  let mockPrismaClient: TransactionClient;

  beforeEach(() => {
    // Create a mock transaction client (excludes $connect, $disconnect, etc.)
    mockPrismaClient = {
      $executeRaw: vi.fn(),
    } as unknown as TransactionClient;
  });

  describe("getRlsClient", () => {
    it("should return undefined when no RLS context is set", () => {
      const client = getRlsClient();
      expect(client).toBeUndefined();
    });

    it("should return the RLS client when context is set", () => {
      runWithRlsClient(mockPrismaClient, () => {
        const client = getRlsClient();
        expect(client).toBe(mockPrismaClient);
      });
    });

    it("should return undefined after context is cleared", () => {
      runWithRlsClient(mockPrismaClient, () => {
        const client = getRlsClient();
        expect(client).toBe(mockPrismaClient);
      });

      // After runWithRlsClient completes, context should be cleared
      const client = getRlsClient();
      expect(client).toBeUndefined();
    });
  });

  describe("runWithRlsClient", () => {
    it("should execute callback with RLS client available", () => {
      const callback = vi.fn(() => {
        const client = getRlsClient();
        expect(client).toBe(mockPrismaClient);
      });

      runWithRlsClient(mockPrismaClient, callback);

      expect(callback).toHaveBeenCalledTimes(1);
    });

    it("should clear context after callback completes", () => {
      runWithRlsClient(mockPrismaClient, () => {
        // Context is set here
      });

      // Context should be cleared after execution
      const client = getRlsClient();
      expect(client).toBeUndefined();
    });

    it("should clear context even if callback throws", () => {
      const error = new Error("Test error");

      expect(() => {
        runWithRlsClient(mockPrismaClient, () => {
          throw error;
        });
      }).toThrow(error);

      // Context should still be cleared
      const client = getRlsClient();
      expect(client).toBeUndefined();
    });

    it("should support nested contexts", () => {
      const outerClient = mockPrismaClient;
      const innerClient = {
        $executeRaw: vi.fn(),
      } as unknown as TransactionClient;

      runWithRlsClient(outerClient, () => {
        expect(getRlsClient()).toBe(outerClient);

        runWithRlsClient(innerClient, () => {
          expect(getRlsClient()).toBe(innerClient);
        });

        // Should restore outer context
        expect(getRlsClient()).toBe(outerClient);
      });

      // Should clear completely after outer context ends
      expect(getRlsClient()).toBeUndefined();
    });
  });
});