import { describe, it, expect, beforeEach, vi } from "vitest"; import { setCurrentUser, setCurrentWorkspace, setWorkspaceContext, clearCurrentUser, clearWorkspaceContext, withUserContext, withUserTransaction, withWorkspaceContext, withAuth, verifyWorkspaceAccess, withoutRLS, createAuthMiddleware, } from "./db-context"; // Mock PrismaClient function createMockPrismaClient(): Record { const mockTx = { $executeRaw: vi.fn().mockResolvedValue(undefined), workspaceMember: { findUnique: vi.fn(), }, workspace: { findMany: vi.fn(), }, }; return { $executeRaw: vi.fn().mockResolvedValue(undefined), $transaction: vi.fn(async (fn: (tx: unknown) => Promise) => { return fn(mockTx); }), workspaceMember: { findUnique: vi.fn(), }, workspace: { findMany: vi.fn(), }, _mockTx: mockTx, // expose for assertions }; } describe("db-context", () => { describe("setCurrentUser", () => { it("should execute SET LOCAL for user ID", async () => { const mockClient = createMockPrismaClient(); await setCurrentUser("user-123", mockClient as never); expect(mockClient.$executeRaw).toHaveBeenCalled(); }); }); describe("setCurrentWorkspace", () => { it("should execute SET LOCAL for workspace ID", async () => { const mockClient = createMockPrismaClient(); await setCurrentWorkspace("ws-123", mockClient as never); expect(mockClient.$executeRaw).toHaveBeenCalled(); }); }); describe("setWorkspaceContext", () => { it("should execute SET LOCAL for both user and workspace", async () => { const mockClient = createMockPrismaClient(); await setWorkspaceContext("user-123", "ws-123", mockClient as never); expect(mockClient.$executeRaw).toHaveBeenCalledTimes(2); }); }); describe("clearCurrentUser", () => { it("should set user ID to NULL", async () => { const mockClient = createMockPrismaClient(); await clearCurrentUser(mockClient as never); expect(mockClient.$executeRaw).toHaveBeenCalled(); }); }); describe("clearWorkspaceContext", () => { it("should set both user and workspace to NULL", async () => { const mockClient = createMockPrismaClient(); await clearWorkspaceContext(mockClient as never); expect(mockClient.$executeRaw).toHaveBeenCalledTimes(2); }); }); describe("withUserContext", () => { it("should execute function within transaction with user context", async () => { // withUserContext uses a global prisma instance, which is hard to mock // without restructuring. We test the higher-level wrappers via // createAuthMiddleware and withWorkspaceContext which accept a client. expect(withUserContext).toBeDefined(); }); }); describe("withUserTransaction", () => { it("should be a function that wraps execution in a transaction", () => { expect(withUserTransaction).toBeDefined(); expect(typeof withUserTransaction).toBe("function"); }); }); describe("withWorkspaceContext", () => { it("should be a function that provides workspace context", () => { expect(withWorkspaceContext).toBeDefined(); expect(typeof withWorkspaceContext).toBe("function"); }); }); describe("withAuth", () => { it("should return a wrapped handler function", () => { const handler = vi.fn().mockResolvedValue("result"); const wrapped = withAuth(handler); expect(typeof wrapped).toBe("function"); }); }); describe("verifyWorkspaceAccess", () => { it("should be a function", () => { expect(verifyWorkspaceAccess).toBeDefined(); expect(typeof verifyWorkspaceAccess).toBe("function"); }); }); describe("withoutRLS", () => { it("should be a function that bypasses RLS", () => { expect(withoutRLS).toBeDefined(); expect(typeof withoutRLS).toBe("function"); }); }); describe("createAuthMiddleware (SEC-API-27)", () => { let mockClient: ReturnType; beforeEach(() => { mockClient = createMockPrismaClient(); }); it("should throw if userId is not provided", async () => { const middleware = createAuthMiddleware(mockClient as never); const next = vi.fn().mockResolvedValue("result"); await expect(middleware({ ctx: { userId: undefined }, next })).rejects.toThrow( "User not authenticated" ); }); it("should call $transaction on the client (RLS context inside transaction)", async () => { const middleware = createAuthMiddleware(mockClient as never); const next = vi.fn().mockResolvedValue("result"); await middleware({ ctx: { userId: "user-123" }, next }); expect(mockClient.$transaction).toHaveBeenCalledTimes(1); expect(mockClient.$transaction).toHaveBeenCalledWith(expect.any(Function)); }); it("should set RLS context inside the transaction, not on the raw client", async () => { const middleware = createAuthMiddleware(mockClient as never); const next = vi.fn().mockResolvedValue("result"); const mockTx = mockClient._mockTx as Record; await middleware({ ctx: { userId: "user-123" }, next }); // The SET LOCAL should be called on the transaction client (mockTx), // NOT on the raw client. This is the core of SEC-API-27. expect(mockTx.$executeRaw as ReturnType).toHaveBeenCalled(); // The raw client's $executeRaw should NOT have been called directly expect(mockClient.$executeRaw).not.toHaveBeenCalled(); }); it("should call next() inside the transaction boundary", async () => { const callOrder: string[] = []; const mockTx = mockClient._mockTx as Record; (mockTx.$executeRaw as ReturnType).mockImplementation(async () => { callOrder.push("setRLS"); }); const next = vi.fn().mockImplementation(async () => { callOrder.push("next"); return "result"; }); // Override $transaction to track that next() is called INSIDE it (mockClient.$transaction as ReturnType).mockImplementation( async (fn: (tx: unknown) => Promise) => { callOrder.push("txStart"); const result = await fn(mockTx); callOrder.push("txEnd"); return result; } ); const middleware = createAuthMiddleware(mockClient as never); await middleware({ ctx: { userId: "user-123" }, next }); expect(callOrder).toEqual(["txStart", "setRLS", "next", "txEnd"]); }); it("should return the result from next()", async () => { const middleware = createAuthMiddleware(mockClient as never); const next = vi.fn().mockResolvedValue({ data: "test" }); const result = await middleware({ ctx: { userId: "user-123" }, next }); expect(result).toEqual({ data: "test" }); }); it("should propagate errors from next() and roll back transaction", async () => { const middleware = createAuthMiddleware(mockClient as never); const error = new Error("Handler error"); const next = vi.fn().mockRejectedValue(error); await expect(middleware({ ctx: { userId: "user-123" }, next })).rejects.toThrow( "Handler error" ); }); it("should not call next() if authentication fails", async () => { const middleware = createAuthMiddleware(mockClient as never); const next = vi.fn().mockResolvedValue("result"); await expect(middleware({ ctx: { userId: undefined }, next })).rejects.toThrow( "User not authenticated" ); expect(next).not.toHaveBeenCalled(); expect(mockClient.$transaction).not.toHaveBeenCalled(); }); }); });