MediaMetz/gemini-cli
1
0
Fork 0
mirror of https://github.com/google-gemini/gemini-cli.git synced 2026-06-13 04:48:09 -07:00
Code Issues Packages Projects Releases Wiki Activity

test: fix broken tests after McpClient refactor

Browse Source
This commit is contained in:
Akhilesh Kumar
2026-03-13 21:38:52 +00:00
parent 6a40da8d97
commit b49b6907e8
3 changed files with 180 additions and 134 deletions
Download Patch File Download Diff File Expand all files Collapse all files
packages/core/src/config/config.test.ts
+1
View File
@@ -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(),
})),
}));
packages/core/src/tools/mcp-client-manager.test.ts
+2
View File
@@ -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');
packages/core/src/tools/mcp-client.test.ts
+177 -134
View File
@@ -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(