chore(cleanup): Consolidate MockTool definitions (#7228)

packages/a2a-server/src/agent.test.ts

@@ -32,8 +32,8 @@ import {
   assertUniqueFinalEventIsLast,
   assertTaskCreationAndWorkingStatus,
   createStreamMessageRequest,
-  MockTool,
 } from './testing_utils.js';
+import { MockTool } from '@google/gemini-cli-core';
 
 const mockToolConfirmationFn = async () =>
   ({}) as unknown as ToolCallConfirmationDetails;
@@ -189,13 +189,10 @@ describe('E2E Tests', () => {
       yield* [];
     });
 
-    const mockTool = new MockTool(
+    const mockTool = new MockTool({
-      'test-tool',
+      name: 'test-tool',
-      'Test Tool',
+      shouldConfirmExecute: vi.fn(mockToolConfirmationFn),
-      true,
+    });
-      false,
-      mockToolConfirmationFn,
-    );
 
     getToolRegistrySpy.mockReturnValue({
       getAllTools: vi.fn().mockReturnValue([mockTool]),
@@ -284,20 +281,16 @@ describe('E2E Tests', () => {
       yield* [];
     });
 
-    const mockTool1 = new MockTool(
+    const mockTool1 = new MockTool({
-      'test-tool-1',
+      name: 'test-tool-1',
-      'Test Tool 1',
+      displayName: 'Test Tool 1',
-      false,
+      shouldConfirmExecute: vi.fn(mockToolConfirmationFn),
-      false,
+    });
-      mockToolConfirmationFn,
+    const mockTool2 = new MockTool({
-    );
+      name: 'test-tool-2',
-    const mockTool2 = new MockTool(
+      displayName: 'Test Tool 2',
-      'test-tool-2',
+      shouldConfirmExecute: vi.fn(mockToolConfirmationFn),
-      'Test Tool 2',
+    });
-      false,
-      false,
-      mockToolConfirmationFn,
-    );
 
     getToolRegistrySpy.mockReturnValue({
       getAllTools: vi.fn().mockReturnValue([mockTool1, mockTool2]),
@@ -404,13 +397,13 @@ describe('E2E Tests', () => {
       yield* [{ type: 'content', value: 'Tool executed successfully.' }];
     });
 
-    const mockTool = new MockTool(
+    const mockTool = new MockTool({
-      'test-tool-no-approval',
+      name: 'test-tool-no-approval',
-      'Test Tool No Approval',
+      displayName: 'Test Tool No Approval',
-    );
+      execute: vi.fn().mockResolvedValue({
-    mockTool.execute.mockResolvedValue({
+        llmContent: 'Tool executed successfully.',
-      llmContent: 'Tool executed successfully.',
+        returnDisplay: 'Tool executed successfully.',
-      returnDisplay: 'Tool executed successfully.',
+      }),
     });
 
     getToolRegistrySpy.mockReturnValue({
@@ -535,15 +528,13 @@ describe('E2E Tests', () => {
     // Set approval mode to yolo
     getApprovalModeSpy.mockReturnValue(ApprovalMode.YOLO);
 
-    const mockTool = new MockTool(
+    const mockTool = new MockTool({
-      'test-tool-yolo',
+      name: 'test-tool-yolo',
-      'Test Tool YOLO',
+      displayName: 'Test Tool YOLO',
-      false,
+      execute: vi.fn().mockResolvedValue({
-      false,
+        llmContent: 'Tool executed successfully.',
-    );
+        returnDisplay: 'Tool executed successfully.',
-    mockTool.execute.mockResolvedValue({
+      }),
-      llmContent: 'Tool executed successfully.',
-      returnDisplay: 'Tool executed successfully.',
     });
 
     getToolRegistrySpy.mockReturnValue({

packages/a2a-server/src/testing_utils.ts

@@ -9,90 +9,7 @@ import type {
   TaskStatusUpdateEvent,
   SendStreamingMessageSuccessResponse,
 } from '@a2a-js/sdk';
-import {
+import { expect } from 'vitest';
-  BaseDeclarativeTool,
-  BaseToolInvocation,
-  Kind,
-} from '@google/gemini-cli-core';
-import type {
-  ToolCallConfirmationDetails,
-  ToolResult,
-  ToolInvocation,
-} from '@google/gemini-cli-core';
-import { expect, vi } from 'vitest';
-
-export const mockOnUserConfirmForToolConfirmation = vi.fn();
-
-export class MockToolInvocation extends BaseToolInvocation<object, ToolResult> {
-  constructor(
-    private readonly tool: MockTool,
-    params: object,
-  ) {
-    super(params);
-  }
-
-  getDescription(): string {
-    return JSON.stringify(this.params);
-  }
-
-  override shouldConfirmExecute(
-    abortSignal: AbortSignal,
-  ): Promise<ToolCallConfirmationDetails | false> {
-    return this.tool.shouldConfirmExecute(this.params, abortSignal);
-  }
-
-  execute(
-    signal: AbortSignal,
-    updateOutput?: (output: string) => void,
-    terminalColumns?: number,
-    terminalRows?: number,
-  ): Promise<ToolResult> {
-    return this.tool.execute(
-      this.params,
-      signal,
-      updateOutput,
-      terminalColumns,
-      terminalRows,
-    );
-  }
-}
-
-// TODO: dedup with gemini-cli, add shouldConfirmExecute() support in core
-export class MockTool extends BaseDeclarativeTool<object, ToolResult> {
-  constructor(
-    name: string,
-    displayName: string,
-    canUpdateOutput = false,
-    isOutputMarkdown = false,
-    shouldConfirmExecute?: () => Promise<ToolCallConfirmationDetails | false>,
-  ) {
-    super(
-      name,
-      displayName,
-      'A mock tool for testing',
-      Kind.Other,
-      {},
-      isOutputMarkdown,
-      canUpdateOutput,
-    );
-
-    if (shouldConfirmExecute) {
-      this.shouldConfirmExecute.mockImplementation(shouldConfirmExecute);
-    } else {
-      // Default to no confirmation needed
-      this.shouldConfirmExecute.mockResolvedValue(false);
-    }
-  }
-
-  execute = vi.fn();
-  shouldConfirmExecute = vi.fn();
-
-  protected createInvocation(
-    params: object,
-  ): ToolInvocation<object, ToolResult> {
-    return new MockToolInvocation(this, params);
-  }
-}
 
 export function createStreamMessageRequest(
   text: string,

packages/cli/src/ui/hooks/useToolScheduler.test.ts

@@ -22,16 +22,13 @@ import type {
   ToolCallResponseInfo,
   ToolCall, // Import from core
   Status as ToolCallStatusType,
-  ToolInvocation,
   AnyDeclarativeTool,
   AnyToolInvocation,
 } from '@google/gemini-cli-core';
 import {
   ToolConfirmationOutcome,
   ApprovalMode,
-  Kind,
+  MockTool,
-  BaseDeclarativeTool,
-  BaseToolInvocation,
 } from '@google/gemini-cli-core';
 import type { HistoryItemWithoutId, HistoryItemToolGroup } from '../types.js';
 import { ToolCallStatus } from '../types.js';
@@ -64,96 +61,20 @@ const mockConfig = {
   }),
 } as unknown as Config;
 
-class MockToolInvocation extends BaseToolInvocation<object, ToolResult> {
+const mockTool = new MockTool({ name: 'mockTool', displayName: 'Mock Tool' });
-  constructor(
+const mockToolWithLiveOutput = new MockTool({
-    private readonly tool: MockTool,
+  name: 'mockToolWithLiveOutput',
-    params: object,
+  displayName: 'Mock Tool With Live Output',
-  ) {
+  description: 'A mock tool for testing',
-    super(params);
+  params: {},
-  }
+  isOutputMarkdown: true,
-
+  canUpdateOutput: true,
-  getDescription(): string {
+});
-    return JSON.stringify(this.params);
-  }
-
-  override shouldConfirmExecute(
-    abortSignal: AbortSignal,
-  ): Promise<ToolCallConfirmationDetails | false> {
-    return this.tool.shouldConfirmExecute(this.params, abortSignal);
-  }
-
-  execute(
-    signal: AbortSignal,
-    updateOutput?: (output: string) => void,
-    terminalColumns?: number,
-    terminalRows?: number,
-  ): Promise<ToolResult> {
-    return this.tool.execute(
-      this.params,
-      signal,
-      updateOutput,
-      terminalColumns,
-      terminalRows,
-    );
-  }
-}
-
-class MockTool extends BaseDeclarativeTool<object, ToolResult> {
-  constructor(
-    name: string,
-    displayName: string,
-    canUpdateOutput = false,
-    shouldConfirm = false,
-    isOutputMarkdown = false,
-  ) {
-    super(
-      name,
-      displayName,
-      'A mock tool for testing',
-      Kind.Other,
-      {},
-      isOutputMarkdown,
-      canUpdateOutput,
-    );
-    if (shouldConfirm) {
-      this.shouldConfirmExecute.mockImplementation(
-        async (): Promise<ToolCallConfirmationDetails | false> => ({
-          type: 'edit',
-          title: 'Mock Tool Requires Confirmation',
-          onConfirm: mockOnUserConfirmForToolConfirmation,
-          filePath: 'mock',
-          fileName: 'mockToolRequiresConfirmation.ts',
-          fileDiff: 'Mock tool requires confirmation',
-          originalContent: 'Original content',
-          newContent: 'New content',
-        }),
-      );
-    }
-  }
-
-  execute = vi.fn();
-  shouldConfirmExecute = vi.fn();
-
-  protected createInvocation(
-    params: object,
-  ): ToolInvocation<object, ToolResult> {
-    return new MockToolInvocation(this, params);
-  }
-}
-
-const mockTool = new MockTool('mockTool', 'Mock Tool');
-const mockToolWithLiveOutput = new MockTool(
-  'mockToolWithLiveOutput',
-  'Mock Tool With Live Output',
-  true,
-);
 let mockOnUserConfirmForToolConfirmation: Mock;
-const mockToolRequiresConfirmation = new MockTool(
+const mockToolRequiresConfirmation = new MockTool({
-  'mockToolRequiresConfirmation',
+  name: 'mockToolRequiresConfirmation',
-  'Mock Tool Requires Confirmation',
+  displayName: 'Mock Tool Requires Confirmation',
-  false,
+});
-  true,
-);
 
 describe('useReactToolScheduler in YOLO Mode', () => {
   let onComplete: Mock;
@@ -717,19 +638,23 @@ describe('useReactToolScheduler', () => {
   });
 
   it('should schedule and execute multiple tool calls', async () => {
-    const tool1 = new MockTool('tool1', 'Tool 1');
+    const tool1 = new MockTool({
-    tool1.execute.mockResolvedValue({
+      name: 'tool1',
-      llmContent: 'Output 1',
+      displayName: 'Tool 1',
-      returnDisplay: 'Display 1',
+      execute: vi.fn().mockResolvedValue({
-    } as ToolResult);
+        llmContent: 'Output 1',
-    tool1.shouldConfirmExecute.mockResolvedValue(null);
+        returnDisplay: 'Display 1',
+      } as ToolResult),
+    });
 
-    const tool2 = new MockTool('tool2', 'Tool 2');
+    const tool2 = new MockTool({
-    tool2.execute.mockResolvedValue({
+      name: 'tool2',
-      llmContent: 'Output 2',
+      displayName: 'Tool 2',
-      returnDisplay: 'Display 2',
+      execute: vi.fn().mockResolvedValue({
-    } as ToolResult);
+        llmContent: 'Output 2',
-    tool2.shouldConfirmExecute.mockResolvedValue(null);
+        returnDisplay: 'Display 2',
+      } as ToolResult),
+    });
 
     mockToolRegistry.getTool.mockImplementation((name) => {
       if (name === 'tool1') return tool1;
@@ -870,7 +795,12 @@ describe('mapToDisplay', () => {
     args: { foo: 'bar' },
   } as any;
 
-  const baseTool = new MockTool('testTool', 'Test Tool Display');
+  const baseTool = new MockTool({
+    name: 'testTool',
+    displayName: 'Test Tool Display',
+    execute: vi.fn(),
+    shouldConfirmExecute: vi.fn(),
+  });
 
   const baseResponse: ToolCallResponseInfo = {
     callId: 'testCallId',
@@ -1028,7 +958,7 @@ describe('mapToDisplay', () => {
       expectedStatus: ToolCallStatus.Error,
       expectedResultDisplay: 'Execution failed display',
       expectedName: baseTool.displayName, // Changed from baseTool.name
-      expectedDescription: baseInvocation.getDescription(),
+      expectedDescription: JSON.stringify(baseRequest.args),
     },
     {
       name: 'cancelled',
@@ -1099,13 +1029,13 @@ describe('mapToDisplay', () => {
       invocation: baseTool.build(baseRequest.args),
       response: { ...baseResponse, callId: 'call1' },
     } as ToolCall;
-    const toolForCall2 = new MockTool(
+    const toolForCall2 = new MockTool({
-      baseTool.name,
+      name: baseTool.name,
-      baseTool.displayName,
+      displayName: baseTool.displayName,
-      false,
+      isOutputMarkdown: true,
-      false,
+      execute: vi.fn(),
-      true,
+      shouldConfirmExecute: vi.fn(),
-    );
+    });
     const toolCall2: ToolCall = {
       request: { ...baseRequest, callId: 'call2' },
       status: 'executing',

packages/core/src/index.ts

@@ -106,3 +106,6 @@ export * from './telemetry/index.js';
 export { sessionId } from './utils/session.js';
 export * from './utils/browser.js';
 export { Storage } from './config/storage.js';
+
+// Export test utils
+export * from './test-utils/index.js';

packages/core/src/test-utils/index.ts

@@ -0,0 +1,7 @@
+/**
+ * @license
+ * Copyright 2025 Google LLC
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+export * from './mock-tool.js';

packages/core/src/test-utils/mock-tool.ts

@@ -0,0 +1,110 @@
+/**
+ * @license
+ * Copyright 2025 Google LLC
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { vi } from 'vitest';
+import type {
+  ToolCallConfirmationDetails,
+  ToolInvocation,
+  ToolResult,
+} from '../tools/tools.js';
+import {
+  BaseDeclarativeTool,
+  BaseToolInvocation,
+  Kind,
+} from '../tools/tools.js';
+
+type MockToolOptions = {
+  name: string;
+  displayName?: string;
+  description?: string;
+  canUpdateOutput?: boolean;
+  isOutputMarkdown?: boolean;
+  shouldConfirmExecute?: (
+    ...args: unknown[]
+  ) => Promise<ToolCallConfirmationDetails | false>;
+  execute?: (...args: unknown[]) => Promise<ToolResult>;
+  params?: object;
+};
+
+class MockToolInvocation extends BaseToolInvocation<
+  { [key: string]: unknown },
+  ToolResult
+> {
+  constructor(
+    private readonly tool: MockTool,
+    params: { [key: string]: unknown },
+  ) {
+    super(params);
+  }
+
+  execute(
+    signal: AbortSignal,
+    updateOutput?: (output: string) => void,
+    terminalColumns?: number,
+    terminalRows?: number,
+  ): Promise<ToolResult> {
+    return this.tool.execute(
+      this.params,
+      signal,
+      updateOutput,
+      terminalColumns,
+      terminalRows,
+    );
+  }
+
+  override shouldConfirmExecute(
+    abortSignal: AbortSignal,
+  ): Promise<ToolCallConfirmationDetails | false> {
+    return this.tool.shouldConfirmExecute(this.params, abortSignal);
+  }
+
+  getDescription(): string {
+    return `A mock tool invocation for ${this.tool.name}`;
+  }
+}
+
+/**
+ * A highly configurable mock tool for testing purposes.
+ */
+export class MockTool extends BaseDeclarativeTool<
+  { [key: string]: unknown },
+  ToolResult
+> {
+  execute: (...args: unknown[]) => Promise<ToolResult>;
+  shouldConfirmExecute: (
+    ...args: unknown[]
+  ) => Promise<ToolCallConfirmationDetails | false>;
+
+  constructor(options: MockToolOptions) {
+    super(
+      options.name,
+      options.displayName ?? options.name,
+      options.description ?? options.name,
+      Kind.Other,
+      options.params,
+      options.isOutputMarkdown ?? false,
+      options.canUpdateOutput ?? false,
+    );
+
+    if (options.shouldConfirmExecute) {
+      this.shouldConfirmExecute = options.shouldConfirmExecute;
+    } else {
+      this.shouldConfirmExecute = vi.fn().mockResolvedValue(false);
+    }
+
+    if (options.execute) {
+      this.execute = options.execute;
+    } else {
+      this.execute = vi.fn();
+    }
+  }
+
+  protected createInvocation(params: {
+    [key: string]: unknown;
+  }): ToolInvocation<{ [key: string]: unknown }, ToolResult> {
+    return new MockToolInvocation(this, params);
+  }
+}