import { describe, it, expect, beforeEach, vi } from 'vitest';
import { ToolRegistry } from './tool-registry.js';
import { MCPTool, MCPMethodDefinition } from '../types/mcp.js';
// Mock logger to avoid console output during tests
vi.mock('../utils/logger.js', () => ({
logger: {
debug: vi.fn(),
info: vi.fn(),
warn: vi.fn(),
error: vi.fn(),
},
}));
// Helper to create a mock tool
function createMockTool(overrides: Partial<MCPTool> = {}): MCPTool {
return {
name: 'test-tool',
description: 'A test tool',
version: '1.0.0',
getMethods: () => [
{
name: 'testMethod',
description: 'A test method',
inputSchema: {},
handler: async () => ({ result: 'ok' }),
},
],
initialize: vi.fn().mockResolvedValue(undefined),
healthCheck: vi.fn().mockResolvedValue(true),
...overrides,
};
}
describe('ToolRegistry', () => {
let registry: ToolRegistry;
beforeEach(() => {
registry = new ToolRegistry();
});
describe('register', () => {
it('should register a tool successfully', async () => {
const tool = createMockTool();
await registry.register(tool);
expect(registry.getTool('test-tool')).toBe(tool);
});
it('should call initialize if provided', async () => {
const initialize = vi.fn().mockResolvedValue(undefined);
const tool = createMockTool({ initialize });
await registry.register(tool);
expect(initialize).toHaveBeenCalledTimes(1);
});
it('should register tool without initialize function', async () => {
const tool = createMockTool({ initialize: undefined });
await registry.register(tool);
expect(registry.getTool('test-tool')).toBe(tool);
});
it('should replace existing tool with same name', async () => {
const tool1 = createMockTool({ version: '1.0.0' });
const tool2 = createMockTool({ version: '2.0.0' });
await registry.register(tool1);
await registry.register(tool2);
expect(registry.getTool('test-tool')?.version).toBe('2.0.0');
});
it('should set health status to true on successful registration', async () => {
const tool = createMockTool();
await registry.register(tool);
const status = registry.getToolStatus('test-tool');
expect(status?.healthy).toBe(true);
});
it('should set health status to false and throw on init failure', async () => {
const initialize = vi.fn().mockRejectedValue(new Error('Init failed'));
const tool = createMockTool({ initialize });
await expect(registry.register(tool)).rejects.toThrow('Init failed');
// After failed init, the tool is not fully registered but health status is set
// The tool may not be in the registry, so getToolStatus returns undefined
// This is expected behavior - a failed tool shouldn't appear in the registry
});
});
describe('unregister', () => {
it('should remove registered tool', async () => {
const tool = createMockTool();
await registry.register(tool);
const result = registry.unregister('test-tool');
expect(result).toBe(true);
expect(registry.getTool('test-tool')).toBeUndefined();
});
it('should return false for non-existent tool', () => {
const result = registry.unregister('non-existent');
expect(result).toBe(false);
});
it('should remove health status when unregistering', async () => {
const tool = createMockTool();
await registry.register(tool);
registry.unregister('test-tool');
expect(registry.getToolStatus('test-tool')).toBeUndefined();
});
});
describe('getTool', () => {
it('should return registered tool', async () => {
const tool = createMockTool();
await registry.register(tool);
expect(registry.getTool('test-tool')).toBe(tool);
});
it('should return undefined for non-existent tool', () => {
expect(registry.getTool('non-existent')).toBeUndefined();
});
});
describe('getAllTools', () => {
it('should return empty array when no tools registered', () => {
expect(registry.getAllTools()).toEqual([]);
});
it('should return all registered tools', async () => {
const tool1 = createMockTool({ name: 'tool-1' });
const tool2 = createMockTool({ name: 'tool-2' });
await registry.register(tool1);
await registry.register(tool2);
const tools = registry.getAllTools();
expect(tools).toHaveLength(2);
expect(tools.map(t => t.name)).toContain('tool-1');
expect(tools.map(t => t.name)).toContain('tool-2');
});
});
describe('getHealthyTools', () => {
it('should return only healthy tools', async () => {
const tool1 = createMockTool({ name: 'healthy-tool' });
const tool2 = createMockTool({
name: 'unhealthy-tool',
initialize: vi.fn().mockRejectedValue(new Error('Failed')),
});
await registry.register(tool1);
try {
await registry.register(tool2);
} catch {
// Expected to fail
}
const healthyTools = registry.getHealthyTools();
expect(healthyTools).toHaveLength(1);
expect(healthyTools[0].name).toBe('healthy-tool');
});
});
describe('getToolStatus', () => {
it('should return tool status with all fields', async () => {
const methods: MCPMethodDefinition[] = [
{
name: 'method1',
description: 'Method 1',
inputSchema: {},
handler: async () => ({}),
},
{
name: 'method2',
description: 'Method 2',
inputSchema: {},
handler: async () => ({}),
},
];
const tool = createMockTool({
name: 'status-tool',
version: '2.0.0',
getMethods: () => methods,
});
await registry.register(tool);
const status = registry.getToolStatus('status-tool');
expect(status).toEqual({
name: 'status-tool',
version: '2.0.0',
healthy: true,
methods: ['method1', 'method2'],
});
});
it('should return undefined for non-existent tool', () => {
expect(registry.getToolStatus('non-existent')).toBeUndefined();
});
});
describe('getAllToolStatus', () => {
it('should return empty array when no tools registered', () => {
expect(registry.getAllToolStatus()).toEqual([]);
});
it('should return status for all tools', async () => {
const tool1 = createMockTool({ name: 'tool-1', version: '1.0.0' });
const tool2 = createMockTool({ name: 'tool-2', version: '2.0.0' });
await registry.register(tool1);
await registry.register(tool2);
const statuses = registry.getAllToolStatus();
expect(statuses).toHaveLength(2);
expect(statuses.map(s => s.name)).toContain('tool-1');
expect(statuses.map(s => s.name)).toContain('tool-2');
});
});
describe('setHealthStatus', () => {
it('should update health status for existing tool', async () => {
const tool = createMockTool();
await registry.register(tool);
registry.setHealthStatus('test-tool', false);
expect(registry.getToolStatus('test-tool')?.healthy).toBe(false);
});
it('should not set status for non-existent tool', () => {
registry.setHealthStatus('non-existent', true);
expect(registry.getToolStatus('non-existent')).toBeUndefined();
});
});
describe('performHealthCheck', () => {
it('should run health check on all tools', async () => {
const healthCheck1 = vi.fn().mockResolvedValue(true);
const healthCheck2 = vi.fn().mockResolvedValue(false);
const tool1 = createMockTool({ name: 'tool-1', healthCheck: healthCheck1 });
const tool2 = createMockTool({ name: 'tool-2', healthCheck: healthCheck2 });
await registry.register(tool1);
await registry.register(tool2);
await registry.performHealthCheck();
expect(healthCheck1).toHaveBeenCalled();
expect(healthCheck2).toHaveBeenCalled();
expect(registry.getToolStatus('tool-1')?.healthy).toBe(true);
expect(registry.getToolStatus('tool-2')?.healthy).toBe(false);
});
it('should set health to false if healthCheck throws', async () => {
const healthCheck = vi.fn().mockRejectedValue(new Error('Health check failed'));
const tool = createMockTool({ healthCheck });
await registry.register(tool);
await registry.performHealthCheck();
expect(registry.getToolStatus('test-tool')?.healthy).toBe(false);
});
it('should skip tools without healthCheck function', async () => {
const tool = createMockTool({ healthCheck: undefined });
await registry.register(tool);
await registry.performHealthCheck();
// Should not throw and health status should remain as set during registration
expect(registry.getToolStatus('test-tool')?.healthy).toBe(true);
});
});
});
