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