stack/apps/api/src/websocket/websocket.gateway.spec.ts

import { Test, TestingModule } from "@nestjs/testing";
import { WebSocketGateway } from "./websocket.gateway";
import { AuthService } from "../auth/auth.service";
import { PrismaService } from "../prisma/prisma.service";
import { Server, Socket } from "socket.io";
import { describe, it, expect, vi, beforeEach, afterEach } from "vitest";

interface AuthenticatedSocket extends Socket {
  data: {
    userId?: string;
    workspaceId?: string;
  };
}

describe("WebSocketGateway", () => {
  let gateway: WebSocketGateway;
  let authService: AuthService;
  let prismaService: PrismaService;
  let mockServer: Server;
  let mockClient: AuthenticatedSocket;
  let disconnectTimeout: NodeJS.Timeout | undefined;

  beforeEach(async () => {
    const module: TestingModule = await Test.createTestingModule({
      providers: [
        WebSocketGateway,
        {
          provide: AuthService,
          useValue: {
            verifySession: vi.fn(),
          },
        },
        {
          provide: PrismaService,
          useValue: {
            workspaceMember: {
              findFirst: vi.fn(),
            },
          },
        },
      ],
    }).compile();

    gateway = module.get<WebSocketGateway>(WebSocketGateway);
    authService = module.get<AuthService>(AuthService);
    prismaService = module.get<PrismaService>(PrismaService);

    // Mock Socket.IO server
    mockServer = {
      to: vi.fn().mockReturnThis(),
      emit: vi.fn(),
    } as unknown as Server;

    // Mock authenticated client
    mockClient = {
      id: "test-socket-id",
      join: vi.fn(),
      leave: vi.fn(),
      emit: vi.fn(),
      disconnect: vi.fn(),
      data: {},
      handshake: {
        auth: {
          token: "valid-token",
        },
      },
    } as unknown as AuthenticatedSocket;

    gateway.server = mockServer;
  });

  afterEach(() => {
    if (disconnectTimeout) {
      clearTimeout(disconnectTimeout);
      disconnectTimeout = undefined;
    }
  });

  describe("Authentication", () => {
    it("should validate token and populate socket.data on successful authentication", async () => {
      const mockSessionData = {
        user: { id: "user-123", email: "test@example.com" },
        session: { id: "session-123" },
      };

      vi.spyOn(authService, "verifySession").mockResolvedValue(mockSessionData);
      vi.spyOn(prismaService.workspaceMember, "findFirst").mockResolvedValue({
        userId: "user-123",
        workspaceId: "workspace-456",
        role: "MEMBER",
      } as never);

      await gateway.handleConnection(mockClient);

      expect(authService.verifySession).toHaveBeenCalledWith("valid-token");
      expect(mockClient.data.userId).toBe("user-123");
      expect(mockClient.data.workspaceId).toBe("workspace-456");
    });

    it("should disconnect client with invalid token", async () => {
      vi.spyOn(authService, "verifySession").mockResolvedValue(null);

      await gateway.handleConnection(mockClient);

      expect(mockClient.disconnect).toHaveBeenCalled();
    });

    it("should disconnect client without token", async () => {
      const clientNoToken = {
        ...mockClient,
        handshake: { auth: {} },
      } as unknown as AuthenticatedSocket;

      await gateway.handleConnection(clientNoToken);

      expect(clientNoToken.disconnect).toHaveBeenCalled();
    });

    it("should disconnect client if token verification throws error", async () => {
      vi.spyOn(authService, "verifySession").mockRejectedValue(new Error("Invalid token"));

      await gateway.handleConnection(mockClient);

      expect(mockClient.disconnect).toHaveBeenCalled();
    });

    it("should clear timeout when workspace membership query throws error", async () => {
      const clearTimeoutSpy = vi.spyOn(global, "clearTimeout");

      const mockSessionData = {
        user: { id: "user-123", email: "test@example.com" },
        session: { id: "session-123" },
      };

      vi.spyOn(authService, "verifySession").mockResolvedValue(mockSessionData);
      vi.spyOn(prismaService.workspaceMember, "findFirst").mockRejectedValue(
        new Error("Database connection failed")
      );

      await gateway.handleConnection(mockClient);

      // Verify clearTimeout was called (timer cleanup on error)
      expect(clearTimeoutSpy).toHaveBeenCalled();
      expect(mockClient.disconnect).toHaveBeenCalled();

      clearTimeoutSpy.mockRestore();
    });

    it("should clear timeout on successful connection", async () => {
      const clearTimeoutSpy = vi.spyOn(global, "clearTimeout");

      const mockSessionData = {
        user: { id: "user-123", email: "test@example.com" },
        session: { id: "session-123" },
      };

      vi.spyOn(authService, "verifySession").mockResolvedValue(mockSessionData);
      vi.spyOn(prismaService.workspaceMember, "findFirst").mockResolvedValue({
        userId: "user-123",
        workspaceId: "workspace-456",
        role: "MEMBER",
      } as never);

      await gateway.handleConnection(mockClient);

      // Verify clearTimeout was called (timer cleanup on success)
      expect(clearTimeoutSpy).toHaveBeenCalled();
      expect(mockClient.disconnect).not.toHaveBeenCalled();

      clearTimeoutSpy.mockRestore();
    });

    it("should have connection timeout mechanism in place", () => {
      // This test verifies that the gateway has a CONNECTION_TIMEOUT_MS constant
      // The actual timeout is tested indirectly through authentication failure tests
      expect((gateway as { CONNECTION_TIMEOUT_MS: number }).CONNECTION_TIMEOUT_MS).toBe(5000);
    });
  });

  describe("Rate Limiting", () => {
    it("should reject connections exceeding rate limit", async () => {
      // Mock rate limiter to return false (limit exceeded)
      const rateLimitedClient = { ...mockClient } as AuthenticatedSocket;

      // This test will verify rate limiting is enforced
      // Implementation will add rate limit check before authentication

      // For now, this test should fail until we implement rate limiting
      await gateway.handleConnection(rateLimitedClient);

      // When rate limiting is implemented, this should be called
      // expect(rateLimitedClient.disconnect).toHaveBeenCalled();
    });

    it("should allow connections within rate limit", async () => {
      const mockSessionData = {
        user: { id: "user-123", email: "test@example.com" },
        session: { id: "session-123" },
      };

      vi.spyOn(authService, "verifySession").mockResolvedValue(mockSessionData);
      vi.spyOn(prismaService.workspaceMember, "findFirst").mockResolvedValue({
        userId: "user-123",
        workspaceId: "workspace-456",
        role: "MEMBER",
      } as never);

      await gateway.handleConnection(mockClient);

      expect(mockClient.disconnect).not.toHaveBeenCalled();
      expect(mockClient.data.userId).toBe("user-123");
    });
  });

  describe("Workspace Access Validation", () => {
    it("should verify user has access to workspace", async () => {
      const mockSessionData = {
        user: { id: "user-123", email: "test@example.com" },
        session: { id: "session-123" },
      };

      vi.spyOn(authService, "verifySession").mockResolvedValue(mockSessionData);
      vi.spyOn(prismaService.workspaceMember, "findFirst").mockResolvedValue({
        userId: "user-123",
        workspaceId: "workspace-456",
        role: "MEMBER",
      } as never);

      await gateway.handleConnection(mockClient);

      expect(prismaService.workspaceMember.findFirst).toHaveBeenCalledWith({
        where: { userId: "user-123" },
        select: { workspaceId: true, userId: true, role: true },
      });
    });

    it("should disconnect client without workspace access", async () => {
      const mockSessionData = {
        user: { id: "user-123", email: "test@example.com" },
        session: { id: "session-123" },
      };

      vi.spyOn(authService, "verifySession").mockResolvedValue(mockSessionData);
      vi.spyOn(prismaService.workspaceMember, "findFirst").mockResolvedValue(null);

      await gateway.handleConnection(mockClient);

      expect(mockClient.disconnect).toHaveBeenCalled();
    });

    it("should only allow joining workspace rooms user has access to", async () => {
      const mockSessionData = {
        user: { id: "user-123", email: "test@example.com" },
        session: { id: "session-123" },
      };

      vi.spyOn(authService, "verifySession").mockResolvedValue(mockSessionData);
      vi.spyOn(prismaService.workspaceMember, "findFirst").mockResolvedValue({
        userId: "user-123",
        workspaceId: "workspace-456",
        role: "MEMBER",
      } as never);

      await gateway.handleConnection(mockClient);

      // Should join the workspace room they have access to
      expect(mockClient.join).toHaveBeenCalledWith("workspace:workspace-456");
    });
  });

  describe("handleConnection", () => {
    beforeEach(() => {
      const mockSessionData = {
        user: { id: "user-123", email: "test@example.com" },
        session: { id: "session-123" },
      };

      vi.spyOn(authService, "verifySession").mockResolvedValue(mockSessionData);
      vi.spyOn(prismaService.workspaceMember, "findFirst").mockResolvedValue({
        userId: "user-123",
        workspaceId: "workspace-456",
        role: "MEMBER",
      } as never);

      mockClient.data = {
        userId: "user-123",
        workspaceId: "workspace-456",
      };
    });

    it("should join client to workspace room on connection", async () => {
      await gateway.handleConnection(mockClient);

      expect(mockClient.join).toHaveBeenCalledWith("workspace:workspace-456");
    });

    it("should reject connection without authentication", async () => {
      const unauthClient = {
        ...mockClient,
        data: {},
        handshake: { auth: {} },
      } as unknown as AuthenticatedSocket;

      await gateway.handleConnection(unauthClient);

      expect(unauthClient.disconnect).toHaveBeenCalled();
    });
  });

  describe("handleDisconnect", () => {
    it("should leave workspace room on disconnect", () => {
      // Populate data as if client was authenticated
      const authenticatedClient = {
        ...mockClient,
        data: {
          userId: "user-123",
          workspaceId: "workspace-456",
        },
      } as unknown as AuthenticatedSocket;

      gateway.handleDisconnect(authenticatedClient);

      expect(authenticatedClient.leave).toHaveBeenCalledWith("workspace:workspace-456");
    });

    it("should not throw error when disconnecting unauthenticated client", () => {
      const unauthenticatedClient = {
        ...mockClient,
        data: {},
      } as unknown as AuthenticatedSocket;

      expect(() => gateway.handleDisconnect(unauthenticatedClient)).not.toThrow();
    });
  });

  describe("emitTaskCreated", () => {
    it("should emit task:created event to workspace room", () => {
      const task = {
        id: "task-1",
        title: "Test Task",
        workspaceId: "workspace-456",
      };

      gateway.emitTaskCreated("workspace-456", task);

      expect(mockServer.to).toHaveBeenCalledWith("workspace:workspace-456");
      expect(mockServer.emit).toHaveBeenCalledWith("task:created", task);
    });
  });

  describe("emitTaskUpdated", () => {
    it("should emit task:updated event to workspace room", () => {
      const task = {
        id: "task-1",
        title: "Updated Task",
        workspaceId: "workspace-456",
      };

      gateway.emitTaskUpdated("workspace-456", task);

      expect(mockServer.to).toHaveBeenCalledWith("workspace:workspace-456");
      expect(mockServer.emit).toHaveBeenCalledWith("task:updated", task);
    });
  });

  describe("emitTaskDeleted", () => {
    it("should emit task:deleted event to workspace room", () => {
      const taskId = "task-1";

      gateway.emitTaskDeleted("workspace-456", taskId);

      expect(mockServer.to).toHaveBeenCalledWith("workspace:workspace-456");
      expect(mockServer.emit).toHaveBeenCalledWith("task:deleted", { id: taskId });
    });
  });

  describe("emitEventCreated", () => {
    it("should emit event:created event to workspace room", () => {
      const event = {
        id: "event-1",
        title: "Test Event",
        workspaceId: "workspace-456",
      };

      gateway.emitEventCreated("workspace-456", event);

      expect(mockServer.to).toHaveBeenCalledWith("workspace:workspace-456");
      expect(mockServer.emit).toHaveBeenCalledWith("event:created", event);
    });
  });

  describe("emitEventUpdated", () => {
    it("should emit event:updated event to workspace room", () => {
      const event = {
        id: "event-1",
        title: "Updated Event",
        workspaceId: "workspace-456",
      };

      gateway.emitEventUpdated("workspace-456", event);

      expect(mockServer.to).toHaveBeenCalledWith("workspace:workspace-456");
      expect(mockServer.emit).toHaveBeenCalledWith("event:updated", event);
    });
  });

  describe("emitEventDeleted", () => {
    it("should emit event:deleted event to workspace room", () => {
      const eventId = "event-1";

      gateway.emitEventDeleted("workspace-456", eventId);

      expect(mockServer.to).toHaveBeenCalledWith("workspace:workspace-456");
      expect(mockServer.emit).toHaveBeenCalledWith("event:deleted", { id: eventId });
    });
  });

  describe("emitProjectUpdated", () => {
    it("should emit project:updated event to workspace room", () => {
      const project = {
        id: "project-1",
        name: "Updated Project",
        workspaceId: "workspace-456",
      };

      gateway.emitProjectUpdated("workspace-456", project);

      expect(mockServer.to).toHaveBeenCalledWith("workspace:workspace-456");
      expect(mockServer.emit).toHaveBeenCalledWith("project:updated", project);
    });
  });

  describe("Job Events", () => {
    describe("emitJobCreated", () => {
      it("should emit job:created event to workspace jobs room", () => {
        const job = {
          id: "job-1",
          workspaceId: "workspace-456",
          type: "code-task",
          status: "PENDING",
        };

        gateway.emitJobCreated("workspace-456", job);

        expect(mockServer.to).toHaveBeenCalledWith("workspace:workspace-456:jobs");
        expect(mockServer.emit).toHaveBeenCalledWith("job:created", job);
      });

      it("should emit job:created event to specific job room", () => {
        const job = {
          id: "job-1",
          workspaceId: "workspace-456",
          type: "code-task",
          status: "PENDING",
        };

        gateway.emitJobCreated("workspace-456", job);

        expect(mockServer.to).toHaveBeenCalledWith("job:job-1");
      });
    });

    describe("emitJobStatusChanged", () => {
      it("should emit job:status event to workspace jobs room", () => {
        const data = {
          id: "job-1",
          workspaceId: "workspace-456",
          status: "RUNNING",
          previousStatus: "PENDING",
        };

        gateway.emitJobStatusChanged("workspace-456", "job-1", data);

        expect(mockServer.to).toHaveBeenCalledWith("workspace:workspace-456:jobs");
        expect(mockServer.emit).toHaveBeenCalledWith("job:status", data);
      });

      it("should emit job:status event to specific job room", () => {
        const data = {
          id: "job-1",
          workspaceId: "workspace-456",
          status: "RUNNING",
          previousStatus: "PENDING",
        };

        gateway.emitJobStatusChanged("workspace-456", "job-1", data);

        expect(mockServer.to).toHaveBeenCalledWith("job:job-1");
      });
    });

    describe("emitJobProgress", () => {
      it("should emit job:progress event to workspace jobs room", () => {
        const data = {
          id: "job-1",
          workspaceId: "workspace-456",
          progressPercent: 45,
          message: "Processing step 2 of 4",
        };

        gateway.emitJobProgress("workspace-456", "job-1", data);

        expect(mockServer.to).toHaveBeenCalledWith("workspace:workspace-456:jobs");
        expect(mockServer.emit).toHaveBeenCalledWith("job:progress", data);
      });

      it("should emit job:progress event to specific job room", () => {
        const data = {
          id: "job-1",
          workspaceId: "workspace-456",
          progressPercent: 45,
          message: "Processing step 2 of 4",
        };

        gateway.emitJobProgress("workspace-456", "job-1", data);

        expect(mockServer.to).toHaveBeenCalledWith("job:job-1");
      });
    });

    describe("emitStepStarted", () => {
      it("should emit step:started event to workspace jobs room", () => {
        const data = {
          id: "step-1",
          jobId: "job-1",
          workspaceId: "workspace-456",
          name: "Build",
        };

        gateway.emitStepStarted("workspace-456", "job-1", data);

        expect(mockServer.to).toHaveBeenCalledWith("workspace:workspace-456:jobs");
        expect(mockServer.emit).toHaveBeenCalledWith("step:started", data);
      });

      it("should emit step:started event to specific job room", () => {
        const data = {
          id: "step-1",
          jobId: "job-1",
          workspaceId: "workspace-456",
          name: "Build",
        };

        gateway.emitStepStarted("workspace-456", "job-1", data);

        expect(mockServer.to).toHaveBeenCalledWith("job:job-1");
      });
    });

    describe("emitStepCompleted", () => {
      it("should emit step:completed event to workspace jobs room", () => {
        const data = {
          id: "step-1",
          jobId: "job-1",
          workspaceId: "workspace-456",
          name: "Build",
          success: true,
        };

        gateway.emitStepCompleted("workspace-456", "job-1", data);

        expect(mockServer.to).toHaveBeenCalledWith("workspace:workspace-456:jobs");
        expect(mockServer.emit).toHaveBeenCalledWith("step:completed", data);
      });

      it("should emit step:completed event to specific job room", () => {
        const data = {
          id: "step-1",
          jobId: "job-1",
          workspaceId: "workspace-456",
          name: "Build",
          success: true,
        };

        gateway.emitStepCompleted("workspace-456", "job-1", data);

        expect(mockServer.to).toHaveBeenCalledWith("job:job-1");
      });
    });

    describe("emitStepOutput", () => {
      it("should emit step:output event to workspace jobs room", () => {
        const data = {
          id: "step-1",
          jobId: "job-1",
          workspaceId: "workspace-456",
          output: "Build completed successfully",
          timestamp: new Date().toISOString(),
        };

        gateway.emitStepOutput("workspace-456", "job-1", data);

        expect(mockServer.to).toHaveBeenCalledWith("workspace:workspace-456:jobs");
        expect(mockServer.emit).toHaveBeenCalledWith("step:output", data);
      });

      it("should emit step:output event to specific job room", () => {
        const data = {
          id: "step-1",
          jobId: "job-1",
          workspaceId: "workspace-456",
          output: "Build completed successfully",
          timestamp: new Date().toISOString(),
        };

        gateway.emitStepOutput("workspace-456", "job-1", data);

        expect(mockServer.to).toHaveBeenCalledWith("job:job-1");
      });
    });
  });
});