stack/apps/api/src/common/guards/workspace.guard.spec.ts

import { describe, it, expect, beforeEach, vi } from "vitest";
import { Test, TestingModule } from "@nestjs/testing";
import { ExecutionContext, ForbiddenException, BadRequestException } from "@nestjs/common";
import { WorkspaceGuard } from "./workspace.guard";
import { PrismaService } from "../../prisma/prisma.service";
import * as dbContext from "../../lib/db-context";

// Mock the db-context module
vi.mock("../../lib/db-context", () => ({
  setCurrentUser: vi.fn(),
}));

describe("WorkspaceGuard", () => {
  let guard: WorkspaceGuard;
  let prismaService: PrismaService;

  const mockPrismaService = {
    workspaceMember: {
      findUnique: vi.fn(),
    },
    $executeRaw: vi.fn(),
  };

  beforeEach(async () => {
    const module: TestingModule = await Test.createTestingModule({
      providers: [
        WorkspaceGuard,
        {
          provide: PrismaService,
          useValue: mockPrismaService,
        },
      ],
    }).compile();

    guard = module.get<WorkspaceGuard>(WorkspaceGuard);
    prismaService = module.get<PrismaService>(PrismaService);

    // Clear all mocks
    vi.clearAllMocks();
  });

  const createMockExecutionContext = (
    user: any,
    headers: Record<string, string> = {},
    params: Record<string, string> = {},
    body: Record<string, any> = {}
  ): ExecutionContext => {
    const mockRequest = {
      user,
      headers,
      params,
      body,
    };

    return {
      switchToHttp: () => ({
        getRequest: () => mockRequest,
      }),
    } as ExecutionContext;
  };

  describe("canActivate", () => {
    const userId = "user-123";
    const workspaceId = "workspace-456";

    it("should allow access when user is a workspace member (via header)", async () => {
      const context = createMockExecutionContext(
        { id: userId },
        { "x-workspace-id": workspaceId }
      );

      mockPrismaService.workspaceMember.findUnique.mockResolvedValue({
        workspaceId,
        userId,
        role: "MEMBER",
      });

      const result = await guard.canActivate(context);

      expect(result).toBe(true);
      expect(mockPrismaService.workspaceMember.findUnique).toHaveBeenCalledWith({
        where: {
          workspaceId_userId: {
            workspaceId,
            userId,
          },
        },
      });
      expect(dbContext.setCurrentUser).toHaveBeenCalledWith(userId, prismaService);

      const request = context.switchToHttp().getRequest();
      expect(request.workspace).toEqual({ id: workspaceId });
      expect(request.user.workspaceId).toBe(workspaceId);
    });

    it("should allow access when user is a workspace member (via URL param)", async () => {
      const context = createMockExecutionContext(
        { id: userId },
        {},
        { workspaceId }
      );

      mockPrismaService.workspaceMember.findUnique.mockResolvedValue({
        workspaceId,
        userId,
        role: "ADMIN",
      });

      const result = await guard.canActivate(context);

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

    it("should allow access when user is a workspace member (via body)", async () => {
      const context = createMockExecutionContext(
        { id: userId },
        {},
        {},
        { workspaceId }
      );

      mockPrismaService.workspaceMember.findUnique.mockResolvedValue({
        workspaceId,
        userId,
        role: "OWNER",
      });

      const result = await guard.canActivate(context);

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

    it("should prioritize header over param and body", async () => {
      const headerWorkspaceId = "workspace-header";
      const paramWorkspaceId = "workspace-param";
      const bodyWorkspaceId = "workspace-body";

      const context = createMockExecutionContext(
        { id: userId },
        { "x-workspace-id": headerWorkspaceId },
        { workspaceId: paramWorkspaceId },
        { workspaceId: bodyWorkspaceId }
      );

      mockPrismaService.workspaceMember.findUnique.mockResolvedValue({
        workspaceId: headerWorkspaceId,
        userId,
        role: "MEMBER",
      });

      await guard.canActivate(context);

      expect(mockPrismaService.workspaceMember.findUnique).toHaveBeenCalledWith({
        where: {
          workspaceId_userId: {
            workspaceId: headerWorkspaceId,
            userId,
          },
        },
      });
    });

    it("should throw ForbiddenException when user is not authenticated", async () => {
      const context = createMockExecutionContext(
        null,
        { "x-workspace-id": workspaceId }
      );

      await expect(guard.canActivate(context)).rejects.toThrow(
        ForbiddenException
      );
      await expect(guard.canActivate(context)).rejects.toThrow(
        "User not authenticated"
      );
    });

    it("should throw BadRequestException when workspace ID is missing", async () => {
      const context = createMockExecutionContext({ id: userId });

      await expect(guard.canActivate(context)).rejects.toThrow(
        BadRequestException
      );
      await expect(guard.canActivate(context)).rejects.toThrow(
        "Workspace ID is required"
      );
    });

    it("should throw ForbiddenException when user is not a workspace member", async () => {
      const context = createMockExecutionContext(
        { id: userId },
        { "x-workspace-id": workspaceId }
      );

      mockPrismaService.workspaceMember.findUnique.mockResolvedValue(null);

      await expect(guard.canActivate(context)).rejects.toThrow(
        ForbiddenException
      );
      await expect(guard.canActivate(context)).rejects.toThrow(
        "You do not have access to this workspace"
      );
    });

    it("should handle database errors gracefully", async () => {
      const context = createMockExecutionContext(
        { id: userId },
        { "x-workspace-id": workspaceId }
      );

      mockPrismaService.workspaceMember.findUnique.mockRejectedValue(
        new Error("Database connection failed")
      );

      await expect(guard.canActivate(context)).rejects.toThrow(
        ForbiddenException
      );
    });
  });
});