stack/apps/api/src/prisma/prisma.service.spec.ts

import { describe, it, expect, beforeEach, afterEach, vi } from "vitest";
import { Test, TestingModule } from "@nestjs/testing";
import { ConfigService } from "@nestjs/config";
import { PrismaService } from "./prisma.service";
import { VaultService } from "../vault/vault.service";
import { CryptoService } from "../federation/crypto.service";

describe("PrismaService", () => {
  let service: PrismaService;
  let mockConfigService: Partial<ConfigService>;

  beforeEach(async () => {
    // Mock ConfigService with a valid test encryption key
    mockConfigService = {
      get: vi.fn((key: string) => {
        const config: Record<string, string> = {
          ENCRYPTION_KEY: "0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef",
          OPENBAO_ADDR: "http://localhost:8200",
          OPENBAO_ROLE_ID: "test-role-id",
          OPENBAO_SECRET_ID: "test-secret-id",
        };
        return config[key] || null;
      }),
    };

    const module: TestingModule = await Test.createTestingModule({
      providers: [
        PrismaService,
        {
          provide: ConfigService,
          useValue: mockConfigService,
        },
        VaultService,
        CryptoService,
      ],
    }).compile();

    service = module.get<PrismaService>(PrismaService);
  });

  afterEach(async () => {
    await service.$disconnect();
    vi.clearAllMocks();
  });

  it("should be defined", () => {
    expect(service).toBeDefined();
  });

  describe("onModuleInit", () => {
    it("should connect to the database and register encryption extensions", async () => {
      const connectSpy = vi.spyOn(service, "$connect").mockResolvedValue(undefined);
      // Mock $extends to return a mock client (extensions chain)
      const mockExtendedClient = { account: {}, llmProviderInstance: {} };
      vi.spyOn(service, "$extends").mockReturnValue({
        ...mockExtendedClient,
        $extends: vi.fn().mockReturnValue(mockExtendedClient),
      } as never);

      await service.onModuleInit();

      expect(connectSpy).toHaveBeenCalled();
    });

    it("should throw error if connection fails", async () => {
      const error = new Error("Connection failed");
      vi.spyOn(service, "$connect").mockRejectedValue(error);

      await expect(service.onModuleInit()).rejects.toThrow(error);
    });
  });

  describe("onModuleDestroy", () => {
    it("should disconnect from the database", async () => {
      const disconnectSpy = vi.spyOn(service, "$disconnect").mockResolvedValue(undefined);

      await service.onModuleDestroy();

      expect(disconnectSpy).toHaveBeenCalled();
    });
  });

  describe("isHealthy", () => {
    it("should return true when database is accessible", async () => {
      vi.spyOn(service, "$queryRaw").mockResolvedValue([{ result: 1 }]);

      const result = await service.isHealthy();

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

    it("should return false when database is not accessible", async () => {
      vi.spyOn(service, "$queryRaw").mockRejectedValue(new Error("Database error"));

      const result = await service.isHealthy();

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

  describe("getConnectionInfo", () => {
    it("should return connection info when connected", async () => {
      const mockResult = [
        {
          current_database: "test_db",
          version: "PostgreSQL 17.0 on x86_64-linux",
        },
      ];
      vi.spyOn(service, "$queryRaw").mockResolvedValue(mockResult);

      const result = await service.getConnectionInfo();

      expect(result).toEqual({
        connected: true,
        database: "test_db",
        version: "PostgreSQL",
      });
    });

    it("should return connected false when result is empty", async () => {
      vi.spyOn(service, "$queryRaw").mockResolvedValue([]);

      const result = await service.getConnectionInfo();

      expect(result).toEqual({ connected: false });
    });

    it("should return connected false when query fails", async () => {
      vi.spyOn(service, "$queryRaw").mockRejectedValue(new Error("Query failed"));

      const result = await service.getConnectionInfo();

      expect(result).toEqual({ connected: false });
    });

    it("should handle version without split", async () => {
      const mockResult = [
        {
          current_database: "test_db",
          version: "PostgreSQL",
        },
      ];
      vi.spyOn(service, "$queryRaw").mockResolvedValue(mockResult);

      const result = await service.getConnectionInfo();

      expect(result).toEqual({
        connected: true,
        database: "test_db",
        version: "PostgreSQL",
      });
    });
  });

  describe("setWorkspaceContext", () => {
    it("should set workspace context variables in transaction", async () => {
      const userId = "user-123";
      const workspaceId = "workspace-456";
      vi.spyOn(service, "$executeRaw").mockResolvedValue(0);

      // Mock $transaction to execute the callback with a mock tx client
      const mockTx = {
        $executeRaw: vi.fn().mockResolvedValue(0),
      };
      vi.spyOn(service, "$transaction").mockImplementation(async (fn) => {
        return fn(mockTx as never);
      });

      await service.$transaction(async (tx) => {
        await service.setWorkspaceContext(userId, workspaceId, tx as never);
      });

      expect(mockTx.$executeRaw).toHaveBeenCalledTimes(2);
    });

    it("should work when called outside transaction using default client", async () => {
      const userId = "user-123";
      const workspaceId = "workspace-456";
      const executeRawSpy = vi.spyOn(service, "$executeRaw").mockResolvedValue(0);

      await service.setWorkspaceContext(userId, workspaceId);

      expect(executeRawSpy).toHaveBeenCalledTimes(2);
    });
  });

  describe("withWorkspaceContext", () => {
    it("should execute function with workspace context set", async () => {
      const userId = "user-123";
      const workspaceId = "workspace-456";

      // Mock $transaction to execute the callback with a mock tx client that has $executeRaw
      const mockTx = {
        $executeRaw: vi.fn().mockResolvedValue(0),
      };

      // Mock both methods at the same time to avoid spy issues
      const setContextCalls: [string, string, unknown][] = [];
      service.setWorkspaceContext = vi.fn().mockImplementation((uid, wid, tx) => {
        setContextCalls.push([uid, wid, tx]);
        return Promise.resolve();
      }) as typeof service.setWorkspaceContext;

      service.$transaction = vi.fn().mockImplementation(async (fn) => {
        return fn(mockTx as never);
      }) as typeof service.$transaction;

      const result = await service.withWorkspaceContext(userId, workspaceId, async () => {
        return "test-result";
      });

      expect(result).toBe("test-result");
      expect(setContextCalls).toHaveLength(1);
      expect(setContextCalls[0]).toEqual([userId, workspaceId, mockTx]);
    });

    it("should pass transaction client to callback", async () => {
      const userId = "user-123";
      const workspaceId = "workspace-456";

      // Mock $transaction to execute the callback with a mock tx client
      const mockTx = {
        $executeRaw: vi.fn().mockResolvedValue(0),
      };

      service.setWorkspaceContext = vi
        .fn()
        .mockResolvedValue(undefined) as typeof service.setWorkspaceContext;
      service.$transaction = vi.fn().mockImplementation(async (fn) => {
        return fn(mockTx as never);
      }) as typeof service.$transaction;

      let receivedClient: unknown = null;
      await service.withWorkspaceContext(userId, workspaceId, async (tx) => {
        receivedClient = tx;
        return null;
      });

      expect(receivedClient).toBeDefined();
      expect(receivedClient).toHaveProperty("$executeRaw");
    });

    it("should handle errors from callback", async () => {
      const userId = "user-123";
      const workspaceId = "workspace-456";

      // Mock $transaction to execute the callback with a mock tx client
      const mockTx = {
        $executeRaw: vi.fn().mockResolvedValue(0),
      };

      service.setWorkspaceContext = vi
        .fn()
        .mockResolvedValue(undefined) as typeof service.setWorkspaceContext;
      service.$transaction = vi.fn().mockImplementation(async (fn) => {
        return fn(mockTx as never);
      }) as typeof service.$transaction;

      const error = new Error("Callback error");
      await expect(
        service.withWorkspaceContext(userId, workspaceId, async () => {
          throw error;
        })
      ).rejects.toThrow(error);
    });
  });

  describe("clearWorkspaceContext", () => {
    it("should clear workspace context variables", async () => {
      // Mock $transaction to execute the callback with a mock tx client
      const mockTx = {
        $executeRaw: vi.fn().mockResolvedValue(0),
      };
      vi.spyOn(service, "$transaction").mockImplementation(async (fn) => {
        return fn(mockTx as never);
      });

      await service.$transaction(async (tx) => {
        await service.clearWorkspaceContext(tx as never);
      });

      expect(mockTx.$executeRaw).toHaveBeenCalledTimes(2);
    });
  });
});