stack/apps/api/src/federation/command.controller.spec.ts

/**
 * Command Controller Tests
 */

import { describe, it, expect, beforeEach, vi } from "vitest";
import { Test, TestingModule } from "@nestjs/testing";
import { CommandController } from "./command.controller";
import { CommandService } from "./command.service";
import { AuthGuard } from "../auth/guards/auth.guard";
import { FederationMessageType, FederationMessageStatus } from "@prisma/client";
import type { AuthenticatedRequest } from "../common/types/user.types";
import type { CommandMessage, CommandResponse } from "./types/message.types";

describe("CommandController", () => {
  let controller: CommandController;
  let commandService: CommandService;

  const mockWorkspaceId = "workspace-123";
  const mockUserId = "user-123";

  beforeEach(async () => {
    const module: TestingModule = await Test.createTestingModule({
      controllers: [CommandController],
      providers: [
        {
          provide: CommandService,
          useValue: {
            sendCommand: vi.fn(),
            handleIncomingCommand: vi.fn(),
            getCommandMessages: vi.fn(),
            getCommandMessage: vi.fn(),
          },
        },
      ],
    })
      .overrideGuard(AuthGuard)
      .useValue({ canActivate: () => true })
      .compile();

    controller = module.get<CommandController>(CommandController);
    commandService = module.get<CommandService>(CommandService);
  });

  describe("sendCommand", () => {
    it("should send a command", async () => {
      const req = {
        user: { id: mockUserId, workspaceId: mockWorkspaceId },
      } as AuthenticatedRequest;

      const dto = {
        connectionId: "conn-123",
        commandType: "spawn_agent",
        payload: { agentType: "task_executor" },
      };

      const mockResult = {
        id: "msg-123",
        workspaceId: mockWorkspaceId,
        connectionId: "conn-123",
        messageType: FederationMessageType.COMMAND,
        messageId: "cmd-123",
        commandType: "spawn_agent",
        payload: { agentType: "task_executor" },
        status: FederationMessageStatus.PENDING,
        createdAt: new Date(),
        updatedAt: new Date(),
      };

      vi.spyOn(commandService, "sendCommand").mockResolvedValue(mockResult as never);

      const result = await controller.sendCommand(req, dto);

      expect(result).toEqual(mockResult);
      expect(commandService.sendCommand).toHaveBeenCalledWith(
        mockWorkspaceId,
        "conn-123",
        "spawn_agent",
        { agentType: "task_executor" }
      );
    });

    it("should throw error if workspace ID not found", async () => {
      const req = {
        user: { id: mockUserId },
      } as AuthenticatedRequest;

      const dto = {
        connectionId: "conn-123",
        commandType: "test",
        payload: {},
      };

      await expect(controller.sendCommand(req, dto)).rejects.toThrow(
        "Workspace ID not found in request"
      );
    });
  });

  describe("handleIncomingCommand", () => {
    it("should handle an incoming command", async () => {
      const dto: CommandMessage = {
        messageId: "cmd-123",
        instanceId: "remote-instance",
        commandType: "spawn_agent",
        payload: { agentType: "task_executor" },
        timestamp: Date.now(),
        signature: "signature-123",
      };

      const mockResponse: CommandResponse = {
        messageId: "resp-123",
        correlationId: "cmd-123",
        instanceId: "local-instance",
        success: true,
        data: { result: "success" },
        timestamp: Date.now(),
        signature: "response-signature",
      };

      vi.spyOn(commandService, "handleIncomingCommand").mockResolvedValue(mockResponse);

      const result = await controller.handleIncomingCommand(dto);

      expect(result).toEqual(mockResponse);
      expect(commandService.handleIncomingCommand).toHaveBeenCalledWith(dto);
    });
  });

  describe("getCommands", () => {
    it("should return all commands for workspace", async () => {
      const req = {
        user: { id: mockUserId, workspaceId: mockWorkspaceId },
      } as AuthenticatedRequest;

      const mockCommands = [
        {
          id: "msg-1",
          workspaceId: mockWorkspaceId,
          connectionId: "conn-123",
          messageType: FederationMessageType.COMMAND,
          messageId: "cmd-1",
          commandType: "test",
          payload: {},
          status: FederationMessageStatus.DELIVERED,
          createdAt: new Date(),
          updatedAt: new Date(),
        },
      ];

      vi.spyOn(commandService, "getCommandMessages").mockResolvedValue(mockCommands as never);

      const result = await controller.getCommands(req);

      expect(result).toEqual(mockCommands);
      expect(commandService.getCommandMessages).toHaveBeenCalledWith(mockWorkspaceId, undefined);
    });

    it("should filter commands by status", async () => {
      const req = {
        user: { id: mockUserId, workspaceId: mockWorkspaceId },
      } as AuthenticatedRequest;

      const mockCommands = [
        {
          id: "msg-1",
          workspaceId: mockWorkspaceId,
          connectionId: "conn-123",
          messageType: FederationMessageType.COMMAND,
          messageId: "cmd-1",
          commandType: "test",
          payload: {},
          status: FederationMessageStatus.PENDING,
          createdAt: new Date(),
          updatedAt: new Date(),
        },
      ];

      vi.spyOn(commandService, "getCommandMessages").mockResolvedValue(mockCommands as never);

      await controller.getCommands(req, FederationMessageStatus.PENDING);

      expect(commandService.getCommandMessages).toHaveBeenCalledWith(
        mockWorkspaceId,
        FederationMessageStatus.PENDING
      );
    });

    it("should throw error if workspace ID not found", async () => {
      const req = {
        user: { id: mockUserId },
      } as AuthenticatedRequest;

      await expect(controller.getCommands(req)).rejects.toThrow(
        "Workspace ID not found in request"
      );
    });
  });

  describe("getCommand", () => {
    it("should return a single command", async () => {
      const req = {
        user: { id: mockUserId, workspaceId: mockWorkspaceId },
      } as AuthenticatedRequest;

      const mockCommand = {
        id: "msg-1",
        workspaceId: mockWorkspaceId,
        connectionId: "conn-123",
        messageType: FederationMessageType.COMMAND,
        messageId: "cmd-1",
        commandType: "test",
        payload: { key: "value" },
        status: FederationMessageStatus.DELIVERED,
        createdAt: new Date(),
        updatedAt: new Date(),
      };

      vi.spyOn(commandService, "getCommandMessage").mockResolvedValue(mockCommand as never);

      const result = await controller.getCommand(req, "msg-1");

      expect(result).toEqual(mockCommand);
      expect(commandService.getCommandMessage).toHaveBeenCalledWith(mockWorkspaceId, "msg-1");
    });

    it("should throw error if workspace ID not found", async () => {
      const req = {
        user: { id: mockUserId },
      } as AuthenticatedRequest;

      await expect(controller.getCommand(req, "msg-1")).rejects.toThrow(
        "Workspace ID not found in request"
      );
    });
  });
});