From b49b6907e82ec9730e5ba7910e6f1b84f67da183 Mon Sep 17 00:00:00 2001 From: Akhilesh Kumar Date: Fri, 13 Mar 2026 21:38:52 +0000 Subject: [PATCH] test: fix broken tests after McpClient refactor --- packages/core/src/config/config.test.ts | 1 + .../core/src/tools/mcp-client-manager.test.ts | 2 + packages/core/src/tools/mcp-client.test.ts | 311 ++++++++++-------- 3 files changed, 180 insertions(+), 134 deletions(-) diff --git a/packages/core/src/config/config.test.ts b/packages/core/src/config/config.test.ts index ee5fc69461..57c10e7a09 100644 --- a/packages/core/src/config/config.test.ts +++ b/packages/core/src/config/config.test.ts @@ -98,6 +98,7 @@ vi.mock('../tools/mcp-client-manager.js', () => ({ McpClientManager: vi.fn().mockImplementation(() => ({ startConfiguredMcpServers: vi.fn(), getMcpInstructions: vi.fn().mockReturnValue('MCP Instructions'), + setMainRegistries: vi.fn(), })), })); diff --git a/packages/core/src/tools/mcp-client-manager.test.ts b/packages/core/src/tools/mcp-client-manager.test.ts index 88e86efc92..61b99915c0 100644 --- a/packages/core/src/tools/mcp-client-manager.test.ts +++ b/packages/core/src/tools/mcp-client-manager.test.ts @@ -17,6 +17,8 @@ import { McpClientManager } from './mcp-client-manager.js'; import { McpClient, MCPDiscoveryState } from './mcp-client.js'; import type { ToolRegistry } from './tool-registry.js'; import type { Config, GeminiCLIExtension } from '../config/config.js'; +import type { PromptRegistry } from '../prompts/prompt-registry.js'; +import type { ResourceRegistry } from '../resources/resource-registry.js'; vi.mock('./mcp-client.js', async () => { const originalModule = await vi.importActual('./mcp-client.js'); diff --git a/packages/core/src/tools/mcp-client.test.ts b/packages/core/src/tools/mcp-client.test.ts index 21b5c28615..4a14b671a0 100644 --- a/packages/core/src/tools/mcp-client.test.ts +++ b/packages/core/src/tools/mcp-client.test.ts @@ -4,6 +4,7 @@ * SPDX-License-Identifier: Apache-2.0 */ +/* eslint-disable @typescript-eslint/no-explicit-any */ import * as ClientLib from '@modelcontextprotocol/sdk/client/index.js'; import { SSEClientTransport } from '@modelcontextprotocol/sdk/client/sse.js'; import * as SdkClientStdioLib from '@modelcontextprotocol/sdk/client/stdio.js'; @@ -160,16 +161,17 @@ describe('mcp-client', () => { { command: 'test-command', }, - mockedToolRegistry, - promptRegistry, - resourceRegistry, workspaceContext, MOCK_CONTEXT, false, '0.0.1', ); await client.connect(); - await client.discover(MOCK_CONTEXT); + await client.discoverInto(MOCK_CONTEXT, { + toolRegistry: mockedToolRegistry, + promptRegistry, + resourceRegistry, + }); expect(mockedClient.listTools).toHaveBeenCalledWith( {}, expect.objectContaining({ timeout: 600000, progressReporter: client }), @@ -244,16 +246,17 @@ describe('mcp-client', () => { { command: 'test-command', }, - mockedToolRegistry, - promptRegistry, - resourceRegistry, workspaceContext, MOCK_CONTEXT, false, '0.0.1', ); await client.connect(); - await client.discover(MOCK_CONTEXT); + await client.discoverInto(MOCK_CONTEXT, { + toolRegistry: mockedToolRegistry, + promptRegistry, + resourceRegistry, + }); expect(mockedToolRegistry.registerTool).toHaveBeenCalledTimes(2); expect(consoleWarnSpy).not.toHaveBeenCalled(); consoleWarnSpy.mockRestore(); @@ -296,16 +299,19 @@ describe('mcp-client', () => { { command: 'test-command', }, - mockedToolRegistry, - promptRegistry, - resourceRegistry, workspaceContext, MOCK_CONTEXT, false, '0.0.1', ); await client.connect(); - await expect(client.discover(MOCK_CONTEXT)).rejects.toThrow('Test error'); + await expect( + client.discoverInto(MOCK_CONTEXT, { + toolRegistry: mockedToolRegistry, + promptRegistry, + resourceRegistry, + }), + ).rejects.toThrow('Test error'); expect(MOCK_CONTEXT.emitMcpDiagnostic).toHaveBeenCalledWith( 'error', `Error discovering prompts from test-server: Test error`, @@ -354,18 +360,19 @@ describe('mcp-client', () => { { command: 'test-command', }, - mockedToolRegistry, - promptRegistry, - resourceRegistry, workspaceContext, MOCK_CONTEXT, false, '0.0.1', ); await client.connect(); - await expect(client.discover(MOCK_CONTEXT)).rejects.toThrow( - 'No prompts, tools, or resources found on the server.', - ); + await expect( + client.discoverInto(MOCK_CONTEXT, { + toolRegistry: mockedToolRegistry, + promptRegistry, + resourceRegistry, + }), + ).rejects.toThrow('No prompts, tools, or resources found on the server.'); }); it('should discover tools if server supports them', async () => { @@ -417,16 +424,17 @@ describe('mcp-client', () => { { command: 'test-command', }, - mockedToolRegistry, - promptRegistry, - resourceRegistry, workspaceContext, MOCK_CONTEXT, false, '0.0.1', ); await client.connect(); - await client.discover(MOCK_CONTEXT); + await client.discoverInto(MOCK_CONTEXT, { + toolRegistry: mockedToolRegistry, + promptRegistry, + resourceRegistry, + }); expect(mockedToolRegistry.registerTool).toHaveBeenCalledOnce(); }); @@ -485,9 +493,6 @@ describe('mcp-client', () => { const client = new McpClient( 'test-server', { command: 'test-command' }, - mockedToolRegistry, - promptRegistry, - resourceRegistry, workspaceContext, MOCK_CONTEXT, false, @@ -495,7 +500,11 @@ describe('mcp-client', () => { ); await client.connect(); - await client.discover(mockConfig); + await client.discoverInto(mockConfig, { + toolRegistry: mockedToolRegistry, + promptRegistry, + resourceRegistry, + }); // Verify tool registration expect(mockedToolRegistry.registerTool).toHaveBeenCalledOnce(); @@ -566,9 +575,6 @@ describe('mcp-client', () => { const client = new McpClient( 'test-server', { command: 'test-command' }, - mockedToolRegistry, - promptRegistry, - resourceRegistry, workspaceContext, MOCK_CONTEXT, false, @@ -576,7 +582,11 @@ describe('mcp-client', () => { ); await client.connect(); - await client.discover(mockConfig); + await client.discoverInto(mockConfig, { + toolRegistry: mockedToolRegistry, + promptRegistry, + resourceRegistry, + }); expect(mockedToolRegistry.registerTool).toHaveBeenCalledOnce(); expect(mockPolicyEngine.addRule).not.toHaveBeenCalled(); @@ -644,9 +654,6 @@ describe('mcp-client', () => { const client = new McpClient( 'test-server', { command: 'test-command' }, - mockedToolRegistry, - promptRegistry, - resourceRegistry, workspaceContext, MOCK_CONTEXT, false, @@ -654,7 +661,11 @@ describe('mcp-client', () => { ); await client.connect(); - await client.discover(mockConfig); + await client.discoverInto(mockConfig, { + toolRegistry: mockedToolRegistry, + promptRegistry, + resourceRegistry, + }); expect(mockedToolRegistry.registerTool).toHaveBeenCalledOnce(); @@ -733,16 +744,17 @@ describe('mcp-client', () => { { command: 'test-command', }, - mockedToolRegistry, - promptRegistry, - resourceRegistry, workspaceContext, MOCK_CONTEXT, false, '0.0.1', ); await client.connect(); - await client.discover(MOCK_CONTEXT); + await client.discoverInto(MOCK_CONTEXT, { + toolRegistry: mockedToolRegistry, + promptRegistry, + resourceRegistry, + }); expect(mockedToolRegistry.registerTool).toHaveBeenCalledOnce(); const registeredTool = vi.mocked(mockedToolRegistry.registerTool).mock .calls[0][0]; @@ -818,16 +830,17 @@ describe('mcp-client', () => { { command: 'test-command', }, - mockedToolRegistry, - promptRegistry, - resourceRegistry, workspaceContext, MOCK_CONTEXT, false, '0.0.1', ); await client.connect(); - await client.discover(MOCK_CONTEXT); + await client.discoverInto(MOCK_CONTEXT, { + toolRegistry: mockedToolRegistry, + promptRegistry, + resourceRegistry, + }); expect(resourceRegistry.setResourcesForServer).toHaveBeenCalledWith( 'test-server', [ @@ -907,16 +920,17 @@ describe('mcp-client', () => { { command: 'test-command', }, - mockedToolRegistry, - promptRegistry, - resourceRegistry, workspaceContext, MOCK_CONTEXT, false, '0.0.1', ); await client.connect(); - await client.discover(MOCK_CONTEXT); + await client.discoverInto(MOCK_CONTEXT, { + toolRegistry: mockedToolRegistry, + promptRegistry, + resourceRegistry, + }); expect(mockedClient.setNotificationHandler).toHaveBeenCalledTimes(2); expect(resourceListHandler).toBeDefined(); @@ -996,16 +1010,17 @@ describe('mcp-client', () => { { command: 'test-command', }, - mockedToolRegistry, - promptRegistry, - resourceRegistry, workspaceContext, MOCK_CONTEXT, false, '0.0.1', ); await client.connect(); - await client.discover(MOCK_CONTEXT); + await client.discoverInto(MOCK_CONTEXT, { + toolRegistry: mockedToolRegistry, + promptRegistry, + resourceRegistry, + }); expect(mockedClient.setNotificationHandler).toHaveBeenCalledTimes(2); expect(promptListHandler).toBeDefined(); @@ -1080,16 +1095,17 @@ describe('mcp-client', () => { { command: 'test-command', }, - mockedToolRegistry, - mockedPromptRegistry, - resourceRegistry, workspaceContext, MOCK_CONTEXT, false, '0.0.1', ); await client.connect(); - await client.discover(MOCK_CONTEXT); + await client.discoverInto(MOCK_CONTEXT, { + toolRegistry: mockedToolRegistry, + promptRegistry: mockedPromptRegistry, + resourceRegistry, + }); expect(mockedToolRegistry.registerTool).toHaveBeenCalledOnce(); expect(mockedPromptRegistry.registerPrompt).toHaveBeenCalledOnce(); @@ -1138,17 +1154,6 @@ describe('mcp-client', () => { const client = new McpClient( 'test-server', { command: 'test-command' }, - mockedToolRegistry, - { - getPromptsByServer: vi.fn().mockReturnValue([]), - registerPrompt: vi.fn(), - } as unknown as PromptRegistry, - { - getResourcesByServer: vi.fn().mockReturnValue([]), - registerResource: vi.fn(), - removeResourcesByServer: vi.fn(), - setResourcesForServer: vi.fn(), - } as unknown as ResourceRegistry, workspaceContext, MOCK_CONTEXT, false, @@ -1156,6 +1161,20 @@ describe('mcp-client', () => { ); await client.connect(); + // INJECTED REGISTRIES + (client as any).registeredRegistries?.add({ + toolRegistry: mockedToolRegistry, + promptRegistry: { + getPromptsByServer: vi.fn().mockReturnValue([]), + registerPrompt: vi.fn(), + } as unknown as PromptRegistry, + resourceRegistry: { + getResourcesByServer: vi.fn().mockReturnValue([]), + registerResource: vi.fn(), + removeResourcesByServer: vi.fn(), + setResourcesForServer: vi.fn(), + } as unknown as ResourceRegistry, + }); expect(mockedClient.setNotificationHandler).toHaveBeenCalledWith( ToolListChangedNotificationSchema, @@ -1183,21 +1202,6 @@ describe('mcp-client', () => { const client = new McpClient( 'test-server', { command: 'test-command' }, - { - getToolsByServer: vi.fn().mockReturnValue([]), - registerTool: vi.fn(), - sortTools: vi.fn(), - } as unknown as ToolRegistry, - { - getPromptsByServer: vi.fn().mockReturnValue([]), - registerPrompt: vi.fn(), - } as unknown as PromptRegistry, - { - getResourcesByServer: vi.fn().mockReturnValue([]), - registerResource: vi.fn(), - removeResourcesByServer: vi.fn(), - setResourcesForServer: vi.fn(), - } as unknown as ResourceRegistry, workspaceContext, MOCK_CONTEXT, false, @@ -1205,6 +1209,24 @@ describe('mcp-client', () => { ); await client.connect(); + // INJECTED REGISTRIES + (client as any).registeredRegistries?.add({ + toolRegistry: { + getToolsByServer: vi.fn().mockReturnValue([]), + registerTool: vi.fn(), + sortTools: vi.fn(), + } as unknown as ToolRegistry, + promptRegistry: { + getPromptsByServer: vi.fn().mockReturnValue([]), + registerPrompt: vi.fn(), + } as unknown as PromptRegistry, + resourceRegistry: { + getResourcesByServer: vi.fn().mockReturnValue([]), + registerResource: vi.fn(), + removeResourcesByServer: vi.fn(), + setResourcesForServer: vi.fn(), + } as unknown as ResourceRegistry, + }); // Should be called for ProgressNotificationSchema, even if no other capabilities expect(mockedClient.setNotificationHandler).toHaveBeenCalled(); @@ -1234,21 +1256,6 @@ describe('mcp-client', () => { const client = new McpClient( 'test-server', { command: 'test-command' }, - { - getToolsByServer: vi.fn().mockReturnValue([]), - registerTool: vi.fn(), - sortTools: vi.fn(), - } as unknown as ToolRegistry, - { - getPromptsByServer: vi.fn().mockReturnValue([]), - registerPrompt: vi.fn(), - } as unknown as PromptRegistry, - { - getResourcesByServer: vi.fn().mockReturnValue([]), - registerResource: vi.fn(), - removeResourcesByServer: vi.fn(), - setResourcesForServer: vi.fn(), - } as unknown as ResourceRegistry, workspaceContext, MOCK_CONTEXT, false, @@ -1256,6 +1263,24 @@ describe('mcp-client', () => { ); await client.connect(); + // INJECTED REGISTRIES + (client as any).registeredRegistries?.add({ + toolRegistry: { + getToolsByServer: vi.fn().mockReturnValue([]), + registerTool: vi.fn(), + sortTools: vi.fn(), + } as unknown as ToolRegistry, + promptRegistry: { + getPromptsByServer: vi.fn().mockReturnValue([]), + registerPrompt: vi.fn(), + } as unknown as PromptRegistry, + resourceRegistry: { + getResourcesByServer: vi.fn().mockReturnValue([]), + registerResource: vi.fn(), + removeResourcesByServer: vi.fn(), + setResourcesForServer: vi.fn(), + } as unknown as ResourceRegistry, + }); const toolUpdateCall = mockedClient.setNotificationHandler.mock.calls.find( @@ -1308,12 +1333,6 @@ describe('mcp-client', () => { const client = new McpClient( 'test-server', { command: 'test-command' }, - mockedToolRegistry, - {} as PromptRegistry, - { - removeMcpResourcesByServer: vi.fn(), - registerResource: vi.fn(), - } as unknown as ResourceRegistry, workspaceContext, MOCK_CONTEXT, false, @@ -1323,6 +1342,15 @@ describe('mcp-client', () => { // 1. Connect (sets up listener) await client.connect(); + // INJECTED REGISTRIES + (client as any).registeredRegistries?.add({ + toolRegistry: mockedToolRegistry, + promptRegistry: {} as PromptRegistry, + resourceRegistry: { + removeMcpResourcesByServer: vi.fn(), + registerResource: vi.fn(), + } as unknown as ResourceRegistry, + }); // 2. Extract the callback passed to setNotificationHandler for tools const toolUpdateCall = @@ -1388,9 +1416,6 @@ describe('mcp-client', () => { const client = new McpClient( 'test-server', { command: 'test-command' }, - mockedToolRegistry, - {} as PromptRegistry, - {} as ResourceRegistry, workspaceContext, MOCK_CONTEXT, false, @@ -1398,6 +1423,12 @@ describe('mcp-client', () => { ); await client.connect(); + // INJECTED REGISTRIES + (client as any).registeredRegistries?.add({ + toolRegistry: mockedToolRegistry, + promptRegistry: {} as PromptRegistry, + resourceRegistry: {} as ResourceRegistry, + }); const toolUpdateCall = mockedClient.setNotificationHandler.mock.calls.find( @@ -1463,9 +1494,6 @@ describe('mcp-client', () => { const clientA = new McpClient( 'server-A', { command: 'cmd-a' }, - mockedToolRegistry, - {} as PromptRegistry, - {} as ResourceRegistry, workspaceContext, MOCK_CONTEXT, false, @@ -1476,9 +1504,6 @@ describe('mcp-client', () => { const clientB = new McpClient( 'server-B', { command: 'cmd-b' }, - mockedToolRegistry, - {} as PromptRegistry, - {} as ResourceRegistry, workspaceContext, MOCK_CONTEXT, false, @@ -1487,7 +1512,19 @@ describe('mcp-client', () => { ); await clientA.connect(); + // INJECTED REGISTRIES + (clientA as any).registeredRegistries?.add({ + toolRegistry: mockedToolRegistry, + promptRegistry: {} as PromptRegistry, + resourceRegistry: {} as ResourceRegistry, + }); await clientB.connect(); + // INJECTED REGISTRIES + (clientB as any).registeredRegistries?.add({ + toolRegistry: mockedToolRegistry, + promptRegistry: {} as PromptRegistry, + resourceRegistry: {} as ResourceRegistry, + }); const toolUpdateCallA = mockClientA.setNotificationHandler.mock.calls.find( @@ -1572,18 +1609,6 @@ describe('mcp-client', () => { 'test-server', // Set a very short timeout { command: 'test-command', timeout: 50 }, - mockedToolRegistry, - { - getPromptsByServer: vi.fn().mockReturnValue([]), - registerPrompt: vi.fn(), - removePromptsByServer: vi.fn(), - } as unknown as PromptRegistry, - { - getResourcesByServer: vi.fn().mockReturnValue([]), - registerResource: vi.fn(), - removeResourcesByServer: vi.fn(), - setResourcesForServer: vi.fn(), - } as unknown as ResourceRegistry, workspaceContext, MOCK_CONTEXT, false, @@ -1591,6 +1616,21 @@ describe('mcp-client', () => { ); await client.connect(); + // INJECTED REGISTRIES + (client as any).registeredRegistries?.add({ + toolRegistry: mockedToolRegistry, + promptRegistry: { + getPromptsByServer: vi.fn().mockReturnValue([]), + registerPrompt: vi.fn(), + removePromptsByServer: vi.fn(), + } as unknown as PromptRegistry, + resourceRegistry: { + getResourcesByServer: vi.fn().mockReturnValue([]), + registerResource: vi.fn(), + removeResourcesByServer: vi.fn(), + setResourcesForServer: vi.fn(), + } as unknown as ResourceRegistry, + }); const toolUpdateCall = mockedClient.setNotificationHandler.mock.calls.find( @@ -1648,18 +1688,6 @@ describe('mcp-client', () => { const client = new McpClient( 'test-server', { command: 'test-command' }, - mockedToolRegistry, - { - getPromptsByServer: vi.fn().mockReturnValue([]), - registerPrompt: vi.fn(), - removePromptsByServer: vi.fn(), - } as unknown as PromptRegistry, - { - getResourcesByServer: vi.fn().mockReturnValue([]), - registerResource: vi.fn(), - removeResourcesByServer: vi.fn(), - setResourcesForServer: vi.fn(), - } as unknown as ResourceRegistry, workspaceContext, MOCK_CONTEXT, false, @@ -1668,6 +1696,21 @@ describe('mcp-client', () => { ); await client.connect(); + // INJECTED REGISTRIES + (client as any).registeredRegistries?.add({ + toolRegistry: mockedToolRegistry, + promptRegistry: { + getPromptsByServer: vi.fn().mockReturnValue([]), + registerPrompt: vi.fn(), + removePromptsByServer: vi.fn(), + } as unknown as PromptRegistry, + resourceRegistry: { + getResourcesByServer: vi.fn().mockReturnValue([]), + registerResource: vi.fn(), + removeResourcesByServer: vi.fn(), + setResourcesForServer: vi.fn(), + } as unknown as ResourceRegistry, + }); const toolUpdateCall = mockedClient.setNotificationHandler.mock.calls.find(