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(
-      'test-tool',
-      'Test Tool',
-      true,
-      false,
-      mockToolConfirmationFn,
-    );
+    const mockTool = new MockTool({
+      name: 'test-tool',
+      shouldConfirmExecute: vi.fn(mockToolConfirmationFn),
+    });
 
     getToolRegistrySpy.mockReturnValue({
       getAllTools: vi.fn().mockReturnValue([mockTool]),
@@ -284,20 +281,16 @@ describe('E2E Tests', () => {
       yield* [];
     });
 
-    const mockTool1 = new MockTool(
-      'test-tool-1',
-      'Test Tool 1',
-      false,
-      false,
-      mockToolConfirmationFn,
-    );
-    const mockTool2 = new MockTool(
-      'test-tool-2',
-      'Test Tool 2',
-      false,
-      false,
-      mockToolConfirmationFn,
-    );
+    const mockTool1 = new MockTool({
+      name: 'test-tool-1',
+      displayName: 'Test Tool 1',
+      shouldConfirmExecute: vi.fn(mockToolConfirmationFn),
+    });
+    const mockTool2 = new MockTool({
+      name: 'test-tool-2',
+      displayName: 'Test Tool 2',
+      shouldConfirmExecute: vi.fn(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(
-      'test-tool-no-approval',
-      'Test Tool No Approval',
-    );
-    mockTool.execute.mockResolvedValue({
-      llmContent: 'Tool executed successfully.',
-      returnDisplay: 'Tool executed successfully.',
+    const mockTool = new MockTool({
+      name: 'test-tool-no-approval',
+      displayName: 'Test Tool No Approval',
+      execute: vi.fn().mockResolvedValue({
+        llmContent: '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(
-      'test-tool-yolo',
-      'Test Tool YOLO',
-      false,
-      false,
-    );
-    mockTool.execute.mockResolvedValue({
-      llmContent: 'Tool executed successfully.',
-      returnDisplay: 'Tool executed successfully.',
+    const mockTool = new MockTool({
+      name: 'test-tool-yolo',
+      displayName: 'Test Tool YOLO',
+      execute: vi.fn().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 {
-  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);
-  }
-}
+import { expect } from 'vitest';
 
 export function createStreamMessageRequest(
   text: string,