stack/apps/api/src/mcp/tool-registry.service.spec.ts

import { describe, it, expect, beforeEach } from "vitest";
import { Test, TestingModule } from "@nestjs/testing";
import { ToolRegistryService } from "./tool-registry.service";
import type { McpTool } from "./interfaces";

describe("ToolRegistryService", () => {
  let service: ToolRegistryService;

  const mockTool1: McpTool = {
    name: "test_tool_1",
    description: "Test tool 1",
    inputSchema: {
      type: "object",
      properties: {
        param1: { type: "string" },
      },
    },
    serverId: "server-1",
  };

  const mockTool2: McpTool = {
    name: "test_tool_2",
    description: "Test tool 2",
    inputSchema: {
      type: "object",
      properties: {
        param2: { type: "number" },
      },
    },
    serverId: "server-1",
  };

  const mockTool3: McpTool = {
    name: "test_tool_3",
    description: "Test tool 3",
    inputSchema: {
      type: "object",
      properties: {
        param3: { type: "boolean" },
      },
    },
    serverId: "server-2",
  };

  beforeEach(async () => {
    const module: TestingModule = await Test.createTestingModule({
      providers: [ToolRegistryService],
    }).compile();

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

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

    it("should start with empty registry", () => {
      const tools = service.listTools();
      expect(tools).toHaveLength(0);
    });
  });

  describe("registerTool", () => {
    it("should register a new tool", () => {
      service.registerTool(mockTool1);
      const tool = service.getTool(mockTool1.name);

      expect(tool).toBeDefined();
      expect(tool?.name).toBe(mockTool1.name);
      expect(tool?.description).toBe(mockTool1.description);
    });

    it("should update existing tool on re-registration", () => {
      service.registerTool(mockTool1);

      const updatedTool: McpTool = {
        ...mockTool1,
        description: "Updated description",
      };

      service.registerTool(updatedTool);
      const tool = service.getTool(mockTool1.name);

      expect(tool?.description).toBe("Updated description");
    });

    it("should register multiple tools", () => {
      service.registerTool(mockTool1);
      service.registerTool(mockTool2);
      service.registerTool(mockTool3);

      const tools = service.listTools();
      expect(tools).toHaveLength(3);
    });
  });

  describe("unregisterTool", () => {
    it("should remove a registered tool", () => {
      service.registerTool(mockTool1);
      service.unregisterTool(mockTool1.name);

      const tool = service.getTool(mockTool1.name);
      expect(tool).toBeUndefined();
    });

    it("should not throw error when unregistering non-existent tool", () => {
      expect(() => service.unregisterTool("non-existent")).not.toThrow();
    });

    it("should only remove the specified tool", () => {
      service.registerTool(mockTool1);
      service.registerTool(mockTool2);

      service.unregisterTool(mockTool1.name);

      expect(service.getTool(mockTool1.name)).toBeUndefined();
      expect(service.getTool(mockTool2.name)).toBeDefined();
    });
  });

  describe("getTool", () => {
    it("should return tool by name", () => {
      service.registerTool(mockTool1);
      const tool = service.getTool(mockTool1.name);

      expect(tool).toEqual(mockTool1);
    });

    it("should return undefined for non-existent tool", () => {
      const tool = service.getTool("non-existent");
      expect(tool).toBeUndefined();
    });
  });

  describe("listTools", () => {
    it("should return all registered tools", () => {
      service.registerTool(mockTool1);
      service.registerTool(mockTool2);
      service.registerTool(mockTool3);

      const tools = service.listTools();

      expect(tools).toHaveLength(3);
      expect(tools).toContainEqual(mockTool1);
      expect(tools).toContainEqual(mockTool2);
      expect(tools).toContainEqual(mockTool3);
    });

    it("should return empty array when no tools registered", () => {
      const tools = service.listTools();
      expect(tools).toHaveLength(0);
    });
  });

  describe("listToolsByServer", () => {
    beforeEach(() => {
      service.registerTool(mockTool1);
      service.registerTool(mockTool2);
      service.registerTool(mockTool3);
    });

    it("should return tools for specific server", () => {
      const server1Tools = service.listToolsByServer("server-1");

      expect(server1Tools).toHaveLength(2);
      expect(server1Tools).toContainEqual(mockTool1);
      expect(server1Tools).toContainEqual(mockTool2);
    });

    it("should return empty array for server with no tools", () => {
      const tools = service.listToolsByServer("non-existent-server");
      expect(tools).toHaveLength(0);
    });

    it("should not include tools from other servers", () => {
      const server2Tools = service.listToolsByServer("server-2");

      expect(server2Tools).toHaveLength(1);
      expect(server2Tools).toContainEqual(mockTool3);
      expect(server2Tools).not.toContainEqual(mockTool1);
    });
  });

  describe("clearServerTools", () => {
    beforeEach(() => {
      service.registerTool(mockTool1);
      service.registerTool(mockTool2);
      service.registerTool(mockTool3);
    });

    it("should remove all tools for a server", () => {
      service.clearServerTools("server-1");

      const server1Tools = service.listToolsByServer("server-1");
      expect(server1Tools).toHaveLength(0);
    });

    it("should not affect tools from other servers", () => {
      service.clearServerTools("server-1");

      const server2Tools = service.listToolsByServer("server-2");
      expect(server2Tools).toHaveLength(1);
    });

    it("should not throw error for non-existent server", () => {
      expect(() => service.clearServerTools("non-existent")).not.toThrow();
    });

    it("should allow re-registration after clearing", () => {
      service.clearServerTools("server-1");
      service.registerTool(mockTool1);

      const tool = service.getTool(mockTool1.name);
      expect(tool).toBeDefined();
    });
  });
});