stack/apps/api/src/mcp/mcp.controller.spec.ts

import { describe, it, expect, beforeEach, vi } from "vitest";
import { Test, TestingModule } from "@nestjs/testing";
import { McpController } from "./mcp.controller";
import { McpHubService } from "./mcp-hub.service";
import { ToolRegistryService } from "./tool-registry.service";
import { NotFoundException } from "@nestjs/common";
import type { McpServerConfig, McpTool } from "./interfaces";
import { RegisterServerDto } from "./dto";

describe("McpController", () => {
  let controller: McpController;
  let hubService: McpHubService;
  let toolRegistry: ToolRegistryService;

  const mockServerConfig: McpServerConfig = {
    id: "test-server",
    name: "Test Server",
    description: "Test MCP server",
    command: "node",
    args: ["server.js"],
  };

  const mockTool: McpTool = {
    name: "test_tool",
    description: "Test tool",
    inputSchema: {
      type: "object",
      properties: {
        param: { type: "string" },
      },
    },
    serverId: "test-server",
  };

  beforeEach(async () => {
    const module: TestingModule = await Test.createTestingModule({
      controllers: [McpController],
      providers: [
        {
          provide: McpHubService,
          useValue: {
            listServers: vi.fn(),
            registerServer: vi.fn(),
            startServer: vi.fn(),
            stopServer: vi.fn(),
            unregisterServer: vi.fn(),
            getServerStatus: vi.fn(),
            sendRequest: vi.fn(),
          },
        },
        {
          provide: ToolRegistryService,
          useValue: {
            listTools: vi.fn(),
            getTool: vi.fn(),
          },
        },
      ],
    }).compile();

    controller = module.get<McpController>(McpController);
    hubService = module.get<McpHubService>(McpHubService);
    toolRegistry = module.get<ToolRegistryService>(ToolRegistryService);
  });

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

  describe("listServers", () => {
    it("should return list of all servers", () => {
      const mockServers = [
        {
          config: mockServerConfig,
          status: "running" as const,
        },
      ];

      vi.spyOn(hubService, "listServers").mockReturnValue(mockServers);

      const result = controller.listServers();

      expect(result).toEqual(mockServers);
      expect(hubService.listServers).toHaveBeenCalled();
    });

    it("should return empty array when no servers", () => {
      vi.spyOn(hubService, "listServers").mockReturnValue([]);

      const result = controller.listServers();

      expect(result).toHaveLength(0);
    });
  });

  describe("registerServer", () => {
    it("should register a new server", async () => {
      const dto: RegisterServerDto = {
        id: mockServerConfig.id,
        name: mockServerConfig.name,
        description: mockServerConfig.description,
        command: mockServerConfig.command,
        args: mockServerConfig.args,
      };

      vi.spyOn(hubService, "registerServer").mockResolvedValue(undefined);

      await controller.registerServer(dto);

      expect(hubService.registerServer).toHaveBeenCalledWith(dto);
    });

    it("should handle registration errors", async () => {
      const dto: RegisterServerDto = {
        id: "test",
        name: "Test",
        description: "Test",
        command: "invalid",
      };

      vi.spyOn(hubService, "registerServer").mockRejectedValue(new Error("Registration failed"));

      await expect(controller.registerServer(dto)).rejects.toThrow("Registration failed");
    });
  });

  describe("startServer", () => {
    it("should start a server by id", async () => {
      vi.spyOn(hubService, "startServer").mockResolvedValue(undefined);

      await controller.startServer(mockServerConfig.id);

      expect(hubService.startServer).toHaveBeenCalledWith(mockServerConfig.id);
    });

    it("should handle start errors", async () => {
      vi.spyOn(hubService, "startServer").mockRejectedValue(new Error("Server not found"));

      await expect(controller.startServer("non-existent")).rejects.toThrow("Server not found");
    });
  });

  describe("stopServer", () => {
    it("should stop a server by id", async () => {
      vi.spyOn(hubService, "stopServer").mockResolvedValue(undefined);

      await controller.stopServer(mockServerConfig.id);

      expect(hubService.stopServer).toHaveBeenCalledWith(mockServerConfig.id);
    });

    it("should handle stop errors", async () => {
      vi.spyOn(hubService, "stopServer").mockRejectedValue(new Error("Server not found"));

      await expect(controller.stopServer("non-existent")).rejects.toThrow("Server not found");
    });
  });

  describe("unregisterServer", () => {
    it("should unregister a server by id", async () => {
      vi.spyOn(hubService, "unregisterServer").mockResolvedValue(undefined);
      vi.spyOn(hubService, "getServerStatus").mockReturnValue({
        config: mockServerConfig,
        status: "stopped",
      });

      await controller.unregisterServer(mockServerConfig.id);

      expect(hubService.unregisterServer).toHaveBeenCalledWith(mockServerConfig.id);
    });

    it("should throw error if server not found", async () => {
      vi.spyOn(hubService, "getServerStatus").mockReturnValue(undefined);

      await expect(controller.unregisterServer("non-existent")).rejects.toThrow(NotFoundException);
    });
  });

  describe("listTools", () => {
    it("should return list of all tools", () => {
      const mockTools = [mockTool];
      vi.spyOn(toolRegistry, "listTools").mockReturnValue(mockTools);

      const result = controller.listTools();

      expect(result).toEqual(mockTools);
      expect(toolRegistry.listTools).toHaveBeenCalled();
    });

    it("should return empty array when no tools", () => {
      vi.spyOn(toolRegistry, "listTools").mockReturnValue([]);

      const result = controller.listTools();

      expect(result).toHaveLength(0);
    });
  });

  describe("getTool", () => {
    it("should return tool by name", () => {
      vi.spyOn(toolRegistry, "getTool").mockReturnValue(mockTool);

      const result = controller.getTool(mockTool.name);

      expect(result).toEqual(mockTool);
      expect(toolRegistry.getTool).toHaveBeenCalledWith(mockTool.name);
    });

    it("should throw NotFoundException if tool not found", () => {
      vi.spyOn(toolRegistry, "getTool").mockReturnValue(undefined);

      expect(() => controller.getTool("non-existent")).toThrow(NotFoundException);
    });
  });

  describe("invokeTool", () => {
    it("should invoke tool and return result", async () => {
      const input = { param: "test value" };
      const expectedResponse = {
        jsonrpc: "2.0" as const,
        id: expect.any(Number),
        result: { success: true },
      };

      vi.spyOn(toolRegistry, "getTool").mockReturnValue(mockTool);
      vi.spyOn(hubService, "sendRequest").mockResolvedValue(expectedResponse);

      const result = await controller.invokeTool(mockTool.name, input);

      expect(result).toEqual({ success: true });
      expect(hubService.sendRequest).toHaveBeenCalledWith(mockTool.serverId, {
        jsonrpc: "2.0",
        id: expect.any(Number),
        method: "tools/call",
        params: {
          name: mockTool.name,
          arguments: input,
        },
      });
    });

    it("should throw NotFoundException if tool not found", async () => {
      vi.spyOn(toolRegistry, "getTool").mockReturnValue(undefined);

      await expect(controller.invokeTool("non-existent", {})).rejects.toThrow(NotFoundException);
    });

    it("should throw error if tool invocation fails", async () => {
      const input = { param: "test" };
      const errorResponse = {
        jsonrpc: "2.0" as const,
        id: 1,
        error: {
          code: -32600,
          message: "Invalid request",
        },
      };

      vi.spyOn(toolRegistry, "getTool").mockReturnValue(mockTool);
      vi.spyOn(hubService, "sendRequest").mockResolvedValue(errorResponse);

      await expect(controller.invokeTool(mockTool.name, input)).rejects.toThrow("Invalid request");
    });
  });
});