Revert "chore(cleanup): Consolidate MockTool definitions (#7228)" (#7283)

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,10 +189,13 @@ describe('E2E Tests', () => {
       yield* [];
     });
 
-    const mockTool = new MockTool({
-      name: 'test-tool',
-      shouldConfirmExecute: vi.fn(mockToolConfirmationFn),
-    });
+    const mockTool = new MockTool(
+      'test-tool',
+      'Test Tool',
+      true,
+      false,
+      mockToolConfirmationFn,
+    );
 
     getToolRegistrySpy.mockReturnValue({
       getAllTools: vi.fn().mockReturnValue([mockTool]),
@@ -281,16 +284,20 @@ describe('E2E Tests', () => {
       yield* [];
     });
 
-    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),
-    });
+    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,
+    );
 
     getToolRegistrySpy.mockReturnValue({
       getAllTools: vi.fn().mockReturnValue([mockTool1, mockTool2]),
@@ -397,13 +404,13 @@ describe('E2E Tests', () => {
       yield* [{ type: 'content', value: '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.',
-      }),
+    const mockTool = new MockTool(
+      'test-tool-no-approval',
+      'Test Tool No Approval',
+    );
+    mockTool.execute.mockResolvedValue({
+      llmContent: 'Tool executed successfully.',
+      returnDisplay: 'Tool executed successfully.',
     });
 
     getToolRegistrySpy.mockReturnValue({
@@ -528,13 +535,15 @@ describe('E2E Tests', () => {
     // Set approval mode to yolo
     getApprovalModeSpy.mockReturnValue(ApprovalMode.YOLO);
 
-    const mockTool = new MockTool({
-      name: 'test-tool-yolo',
-      displayName: 'Test Tool YOLO',
-      execute: vi.fn().mockResolvedValue({
-        llmContent: 'Tool executed successfully.',
-        returnDisplay: 'Tool executed successfully.',
-      }),
+    const mockTool = new MockTool(
+      'test-tool-yolo',
+      'Test Tool YOLO',
+      false,
+      false,
+    );
+    mockTool.execute.mockResolvedValue({
+      llmContent: 'Tool executed successfully.',
+      returnDisplay: 'Tool executed successfully.',
     });
 
     getToolRegistrySpy.mockReturnValue({

packages/a2a-server/src/testing_utils.ts

@@ -9,7 +9,90 @@ import type {
   TaskStatusUpdateEvent,
   SendStreamingMessageSuccessResponse,
 } from '@a2a-js/sdk';
-import { expect } from 'vitest';
+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);
+  }
+}
 
 export function createStreamMessageRequest(
   text: string,