From f0146c8b85ef68ea3b0cd4bcf61ad54378443a8f Mon Sep 17 00:00:00 2001 From: Adam Weidman <65992621+adamfweidman@users.noreply.github.com> Date: Wed, 27 Aug 2025 18:18:26 +0000 Subject: [PATCH] chore(cleanup): Consolidate MockTool definitions (#7228) --- packages/a2a-server/src/agent.test.ts | 67 ++++---- packages/a2a-server/src/testing_utils.ts | 85 +--------- .../cli/src/ui/hooks/useToolScheduler.test.ts | 158 +++++------------- packages/core/src/index.ts | 3 + packages/core/src/test-utils/index.ts | 7 + packages/core/src/test-utils/mock-tool.ts | 110 ++++++++++++ 6 files changed, 194 insertions(+), 236 deletions(-) create mode 100644 packages/core/src/test-utils/index.ts create mode 100644 packages/core/src/test-utils/mock-tool.ts diff --git a/packages/a2a-server/src/agent.test.ts b/packages/a2a-server/src/agent.test.ts index 04160d3f..4c563179 100644 --- a/packages/a2a-server/src/agent.test.ts +++ b/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({ diff --git a/packages/a2a-server/src/testing_utils.ts b/packages/a2a-server/src/testing_utils.ts index bd7ddaaa..b9c1de17 100644 --- a/packages/a2a-server/src/testing_utils.ts +++ b/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 { - constructor( - private readonly tool: MockTool, - params: object, - ) { - super(params); - } - - getDescription(): string { - return JSON.stringify(this.params); - } - - override shouldConfirmExecute( - abortSignal: AbortSignal, - ): Promise { - return this.tool.shouldConfirmExecute(this.params, abortSignal); - } - - execute( - signal: AbortSignal, - updateOutput?: (output: string) => void, - terminalColumns?: number, - terminalRows?: number, - ): Promise { - 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 { - constructor( - name: string, - displayName: string, - canUpdateOutput = false, - isOutputMarkdown = false, - shouldConfirmExecute?: () => Promise, - ) { - 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 { - return new MockToolInvocation(this, params); - } -} +import { expect } from 'vitest'; export function createStreamMessageRequest( text: string, diff --git a/packages/cli/src/ui/hooks/useToolScheduler.test.ts b/packages/cli/src/ui/hooks/useToolScheduler.test.ts index 0100b551..4584cfa3 100644 --- a/packages/cli/src/ui/hooks/useToolScheduler.test.ts +++ b/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, - BaseDeclarativeTool, - BaseToolInvocation, + MockTool, } 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 { - constructor( - private readonly tool: MockTool, - params: object, - ) { - super(params); - } - - getDescription(): string { - return JSON.stringify(this.params); - } - - override shouldConfirmExecute( - abortSignal: AbortSignal, - ): Promise { - return this.tool.shouldConfirmExecute(this.params, abortSignal); - } - - execute( - signal: AbortSignal, - updateOutput?: (output: string) => void, - terminalColumns?: number, - terminalRows?: number, - ): Promise { - return this.tool.execute( - this.params, - signal, - updateOutput, - terminalColumns, - terminalRows, - ); - } -} - -class MockTool extends BaseDeclarativeTool { - 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 => ({ - 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 { - return new MockToolInvocation(this, params); - } -} - -const mockTool = new MockTool('mockTool', 'Mock Tool'); -const mockToolWithLiveOutput = new MockTool( - 'mockToolWithLiveOutput', - 'Mock Tool With Live Output', - true, -); +const mockTool = new MockTool({ name: 'mockTool', displayName: 'Mock Tool' }); +const mockToolWithLiveOutput = new MockTool({ + name: 'mockToolWithLiveOutput', + displayName: 'Mock Tool With Live Output', + description: 'A mock tool for testing', + params: {}, + isOutputMarkdown: true, + canUpdateOutput: true, +}); let mockOnUserConfirmForToolConfirmation: Mock; -const mockToolRequiresConfirmation = new MockTool( - 'mockToolRequiresConfirmation', - 'Mock Tool Requires Confirmation', - false, - true, -); +const mockToolRequiresConfirmation = new MockTool({ + name: 'mockToolRequiresConfirmation', + displayName: 'Mock Tool Requires Confirmation', +}); 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'); - tool1.execute.mockResolvedValue({ - llmContent: 'Output 1', - returnDisplay: 'Display 1', - } as ToolResult); - tool1.shouldConfirmExecute.mockResolvedValue(null); + const tool1 = new MockTool({ + name: 'tool1', + displayName: 'Tool 1', + execute: vi.fn().mockResolvedValue({ + llmContent: 'Output 1', + returnDisplay: 'Display 1', + } as ToolResult), + }); - const tool2 = new MockTool('tool2', 'Tool 2'); - tool2.execute.mockResolvedValue({ - llmContent: 'Output 2', - returnDisplay: 'Display 2', - } as ToolResult); - tool2.shouldConfirmExecute.mockResolvedValue(null); + const tool2 = new MockTool({ + name: 'tool2', + displayName: 'Tool 2', + execute: vi.fn().mockResolvedValue({ + llmContent: 'Output 2', + 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( - baseTool.name, - baseTool.displayName, - false, - false, - true, - ); + const toolForCall2 = new MockTool({ + name: baseTool.name, + displayName: baseTool.displayName, + isOutputMarkdown: true, + execute: vi.fn(), + shouldConfirmExecute: vi.fn(), + }); const toolCall2: ToolCall = { request: { ...baseRequest, callId: 'call2' }, status: 'executing', diff --git a/packages/core/src/index.ts b/packages/core/src/index.ts index 64732c21..c8a2dda1 100644 --- a/packages/core/src/index.ts +++ b/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'; diff --git a/packages/core/src/test-utils/index.ts b/packages/core/src/test-utils/index.ts new file mode 100644 index 00000000..6146d39d --- /dev/null +++ b/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'; diff --git a/packages/core/src/test-utils/mock-tool.ts b/packages/core/src/test-utils/mock-tool.ts new file mode 100644 index 00000000..5fb6a6d2 --- /dev/null +++ b/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; + execute?: (...args: unknown[]) => Promise; + 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 { + return this.tool.execute( + this.params, + signal, + updateOutput, + terminalColumns, + terminalRows, + ); + } + + override shouldConfirmExecute( + abortSignal: AbortSignal, + ): Promise { + 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; + shouldConfirmExecute: ( + ...args: unknown[] + ) => Promise; + + 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); + } +}