summaryrefslogtreecommitdiff
path: root/packages/core/src/tools/write-file.test.ts
diff options
context:
space:
mode:
authorTommaso Sciortino <[email protected]>2025-05-30 18:25:47 -0700
committerGitHub <[email protected]>2025-05-30 18:25:47 -0700
commit21fba832d1b4ea7af43fb887d9b2b38fcf8210d0 (patch)
tree7200d2fac3a55c385e0a2dac34b5282c942364bc /packages/core/src/tools/write-file.test.ts
parentc81148a0cc8489f657901c2cc7247c0834075e1a (diff)
Rename server->core (#638)
Diffstat (limited to 'packages/core/src/tools/write-file.test.ts')
-rw-r--r--packages/core/src/tools/write-file.test.ts567
1 files changed, 567 insertions, 0 deletions
diff --git a/packages/core/src/tools/write-file.test.ts b/packages/core/src/tools/write-file.test.ts
new file mode 100644
index 00000000..3fd97c9e
--- /dev/null
+++ b/packages/core/src/tools/write-file.test.ts
@@ -0,0 +1,567 @@
+/**
+ * @license
+ * Copyright 2025 Google LLC
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import {
+ describe,
+ it,
+ expect,
+ beforeEach,
+ afterEach,
+ vi,
+ type Mocked,
+} from 'vitest';
+import { WriteFileTool } from './write-file.js';
+import {
+ FileDiff,
+ ToolConfirmationOutcome,
+ ToolEditConfirmationDetails,
+} from './tools.js';
+import { type EditToolParams } from './edit.js';
+import { Config } from '../config/config.js';
+import { ToolRegistry } from './tool-registry.js';
+import path from 'path';
+import fs from 'fs';
+import os from 'os';
+import { GeminiClient } from '../core/client.js';
+import {
+ ensureCorrectEdit,
+ ensureCorrectFileContent,
+ CorrectedEditResult,
+} from '../utils/editCorrector.js';
+
+const rootDir = path.resolve(os.tmpdir(), 'gemini-cli-test-root');
+
+// --- MOCKS ---
+vi.mock('../core/client.js');
+vi.mock('../utils/editCorrector.js');
+
+let mockGeminiClientInstance: Mocked<GeminiClient>;
+const mockEnsureCorrectEdit = vi.fn<typeof ensureCorrectEdit>();
+const mockEnsureCorrectFileContent = vi.fn<typeof ensureCorrectFileContent>();
+
+// Wire up the mocked functions to be used by the actual module imports
+vi.mocked(ensureCorrectEdit).mockImplementation(mockEnsureCorrectEdit);
+vi.mocked(ensureCorrectFileContent).mockImplementation(
+ mockEnsureCorrectFileContent,
+);
+
+// Mock Config
+const mockConfigInternal = {
+ getTargetDir: () => rootDir,
+ getAlwaysSkipModificationConfirmation: vi.fn(() => false),
+ setAlwaysSkipModificationConfirmation: vi.fn(),
+ getApiKey: () => 'test-key',
+ getModel: () => 'test-model',
+ getSandbox: () => false,
+ getDebugMode: () => false,
+ getQuestion: () => undefined,
+ getFullContext: () => false,
+ getToolDiscoveryCommand: () => undefined,
+ getToolCallCommand: () => undefined,
+ getMcpServerCommand: () => undefined,
+ getMcpServers: () => undefined,
+ getUserAgent: () => 'test-agent',
+ getUserMemory: () => '',
+ setUserMemory: vi.fn(),
+ getGeminiMdFileCount: () => 0,
+ setGeminiMdFileCount: vi.fn(),
+ getToolRegistry: () =>
+ ({
+ registerTool: vi.fn(),
+ discoverTools: vi.fn(),
+ }) as unknown as ToolRegistry,
+};
+const mockConfig = mockConfigInternal as unknown as Config;
+// --- END MOCKS ---
+
+describe('WriteFileTool', () => {
+ let tool: WriteFileTool;
+ let tempDir: string;
+
+ beforeEach(() => {
+ // Create a unique temporary directory for files created outside the root
+ tempDir = fs.mkdtempSync(
+ path.join(os.tmpdir(), 'write-file-test-external-'),
+ );
+ // Ensure the rootDir for the tool exists
+ if (!fs.existsSync(rootDir)) {
+ fs.mkdirSync(rootDir, { recursive: true });
+ }
+
+ // Setup GeminiClient mock
+ mockGeminiClientInstance = new (vi.mocked(GeminiClient))(
+ mockConfig,
+ ) as Mocked<GeminiClient>;
+ vi.mocked(GeminiClient).mockImplementation(() => mockGeminiClientInstance);
+
+ tool = new WriteFileTool(mockConfig);
+
+ // Reset mocks before each test
+ mockConfigInternal.getAlwaysSkipModificationConfirmation.mockReturnValue(
+ false,
+ );
+ mockConfigInternal.setAlwaysSkipModificationConfirmation.mockClear();
+ mockEnsureCorrectEdit.mockReset();
+ mockEnsureCorrectFileContent.mockReset();
+
+ // Default mock implementations that return valid structures
+ mockEnsureCorrectEdit.mockImplementation(
+ async (
+ _currentContent: string,
+ params: EditToolParams,
+ _client: GeminiClient,
+ signal?: AbortSignal, // Make AbortSignal optional to match usage
+ ): Promise<CorrectedEditResult> => {
+ if (signal?.aborted) {
+ return Promise.reject(new Error('Aborted'));
+ }
+ return Promise.resolve({
+ params: { ...params, new_string: params.new_string ?? '' },
+ occurrences: 1,
+ });
+ },
+ );
+ mockEnsureCorrectFileContent.mockImplementation(
+ async (
+ content: string,
+ _client: GeminiClient,
+ signal?: AbortSignal,
+ ): Promise<string> => {
+ // Make AbortSignal optional
+ if (signal?.aborted) {
+ return Promise.reject(new Error('Aborted'));
+ }
+ return Promise.resolve(content ?? '');
+ },
+ );
+ });
+
+ afterEach(() => {
+ // Clean up the temporary directories
+ if (fs.existsSync(tempDir)) {
+ fs.rmSync(tempDir, { recursive: true, force: true });
+ }
+ if (fs.existsSync(rootDir)) {
+ fs.rmSync(rootDir, { recursive: true, force: true });
+ }
+ vi.clearAllMocks();
+ });
+
+ describe('validateToolParams', () => {
+ it('should return null for valid absolute path within root', () => {
+ const params = {
+ file_path: path.join(rootDir, 'test.txt'),
+ content: 'hello',
+ };
+ expect(tool.validateToolParams(params)).toBeNull();
+ });
+
+ it('should return error for relative path', () => {
+ const params = { file_path: 'test.txt', content: 'hello' };
+ expect(tool.validateToolParams(params)).toMatch(
+ /File path must be absolute/,
+ );
+ });
+
+ it('should return error for path outside root', () => {
+ const outsidePath = path.resolve(tempDir, 'outside-root.txt');
+ const params = {
+ file_path: outsidePath,
+ content: 'hello',
+ };
+ expect(tool.validateToolParams(params)).toMatch(
+ /File path must be within the root directory/,
+ );
+ });
+
+ it('should return error if path is a directory', () => {
+ const dirAsFilePath = path.join(rootDir, 'a_directory');
+ fs.mkdirSync(dirAsFilePath);
+ const params = {
+ file_path: dirAsFilePath,
+ content: 'hello',
+ };
+ expect(tool.validateToolParams(params)).toMatch(
+ `Path is a directory, not a file: ${dirAsFilePath}`,
+ );
+ });
+ });
+
+ describe('_getCorrectedFileContent', () => {
+ it('should call ensureCorrectFileContent for a new file', async () => {
+ const filePath = path.join(rootDir, 'new_corrected_file.txt');
+ const proposedContent = 'Proposed new content.';
+ const correctedContent = 'Corrected new content.';
+ const abortSignal = new AbortController().signal;
+ // Ensure the mock is set for this specific test case if needed, or rely on beforeEach
+ mockEnsureCorrectFileContent.mockResolvedValue(correctedContent);
+
+ // @ts-expect-error _getCorrectedFileContent is private
+ const result = await tool._getCorrectedFileContent(
+ filePath,
+ proposedContent,
+ abortSignal,
+ );
+
+ expect(mockEnsureCorrectFileContent).toHaveBeenCalledWith(
+ proposedContent,
+ mockGeminiClientInstance,
+ abortSignal,
+ );
+ expect(mockEnsureCorrectEdit).not.toHaveBeenCalled();
+ expect(result.correctedContent).toBe(correctedContent);
+ expect(result.originalContent).toBe('');
+ expect(result.fileExists).toBe(false);
+ expect(result.error).toBeUndefined();
+ });
+
+ it('should call ensureCorrectEdit for an existing file', async () => {
+ const filePath = path.join(rootDir, 'existing_corrected_file.txt');
+ const originalContent = 'Original existing content.';
+ const proposedContent = 'Proposed replacement content.';
+ const correctedProposedContent = 'Corrected replacement content.';
+ const abortSignal = new AbortController().signal;
+ fs.writeFileSync(filePath, originalContent, 'utf8');
+
+ // Ensure this mock is active and returns the correct structure
+ mockEnsureCorrectEdit.mockResolvedValue({
+ params: {
+ file_path: filePath,
+ old_string: originalContent,
+ new_string: correctedProposedContent,
+ },
+ occurrences: 1,
+ } as CorrectedEditResult);
+
+ // @ts-expect-error _getCorrectedFileContent is private
+ const result = await tool._getCorrectedFileContent(
+ filePath,
+ proposedContent,
+ abortSignal,
+ );
+
+ expect(mockEnsureCorrectEdit).toHaveBeenCalledWith(
+ originalContent,
+ {
+ old_string: originalContent,
+ new_string: proposedContent,
+ file_path: filePath,
+ },
+ mockGeminiClientInstance,
+ abortSignal,
+ );
+ expect(mockEnsureCorrectFileContent).not.toHaveBeenCalled();
+ expect(result.correctedContent).toBe(correctedProposedContent);
+ expect(result.originalContent).toBe(originalContent);
+ expect(result.fileExists).toBe(true);
+ expect(result.error).toBeUndefined();
+ });
+
+ it('should return error if reading an existing file fails (e.g. permissions)', async () => {
+ const filePath = path.join(rootDir, 'unreadable_file.txt');
+ const proposedContent = 'some content';
+ const abortSignal = new AbortController().signal;
+ fs.writeFileSync(filePath, 'content', { mode: 0o000 });
+
+ const readError = new Error('Permission denied');
+ const originalReadFileSync = fs.readFileSync;
+ vi.spyOn(fs, 'readFileSync').mockImplementationOnce(() => {
+ throw readError;
+ });
+
+ // @ts-expect-error _getCorrectedFileContent is private
+ const result = await tool._getCorrectedFileContent(
+ filePath,
+ proposedContent,
+ abortSignal,
+ );
+
+ expect(fs.readFileSync).toHaveBeenCalledWith(filePath, 'utf8');
+ expect(mockEnsureCorrectEdit).not.toHaveBeenCalled();
+ expect(mockEnsureCorrectFileContent).not.toHaveBeenCalled();
+ expect(result.correctedContent).toBe(proposedContent);
+ expect(result.originalContent).toBe('');
+ expect(result.fileExists).toBe(true);
+ expect(result.error).toEqual({
+ message: 'Permission denied',
+ code: undefined,
+ });
+
+ vi.spyOn(fs, 'readFileSync').mockImplementation(originalReadFileSync);
+ fs.chmodSync(filePath, 0o600);
+ });
+ });
+
+ describe('shouldConfirmExecute', () => {
+ const abortSignal = new AbortController().signal;
+ it('should return false if params are invalid (relative path)', async () => {
+ const params = { file_path: 'relative.txt', content: 'test' };
+ const confirmation = await tool.shouldConfirmExecute(params, abortSignal);
+ expect(confirmation).toBe(false);
+ });
+
+ it('should return false if params are invalid (outside root)', async () => {
+ const outsidePath = path.resolve(tempDir, 'outside-root.txt');
+ const params = { file_path: outsidePath, content: 'test' };
+ const confirmation = await tool.shouldConfirmExecute(params, abortSignal);
+ expect(confirmation).toBe(false);
+ });
+
+ it('should return false if _getCorrectedFileContent returns an error', async () => {
+ const filePath = path.join(rootDir, 'confirm_error_file.txt');
+ const params = { file_path: filePath, content: 'test content' };
+ fs.writeFileSync(filePath, 'original', { mode: 0o000 });
+
+ const readError = new Error('Simulated read error for confirmation');
+ const originalReadFileSync = fs.readFileSync;
+ vi.spyOn(fs, 'readFileSync').mockImplementationOnce(() => {
+ throw readError;
+ });
+
+ const confirmation = await tool.shouldConfirmExecute(params, abortSignal);
+ expect(confirmation).toBe(false);
+
+ vi.spyOn(fs, 'readFileSync').mockImplementation(originalReadFileSync);
+ fs.chmodSync(filePath, 0o600);
+ });
+
+ it('should request confirmation with diff for a new file (with corrected content)', async () => {
+ const filePath = path.join(rootDir, 'confirm_new_file.txt');
+ const proposedContent = 'Proposed new content for confirmation.';
+ const correctedContent = 'Corrected new content for confirmation.';
+ mockEnsureCorrectFileContent.mockResolvedValue(correctedContent); // Ensure this mock is active
+
+ const params = { file_path: filePath, content: proposedContent };
+ const confirmation = (await tool.shouldConfirmExecute(
+ params,
+ abortSignal,
+ )) as ToolEditConfirmationDetails;
+
+ expect(mockEnsureCorrectFileContent).toHaveBeenCalledWith(
+ proposedContent,
+ mockGeminiClientInstance,
+ abortSignal,
+ );
+ expect(confirmation).toEqual(
+ expect.objectContaining({
+ title: `Confirm Write: ${path.basename(filePath)}`,
+ fileName: 'confirm_new_file.txt',
+ fileDiff: expect.stringContaining(correctedContent),
+ }),
+ );
+ expect(confirmation.fileDiff).toMatch(
+ /--- confirm_new_file.txt\tCurrent/,
+ );
+ expect(confirmation.fileDiff).toMatch(
+ /\+\+\+ confirm_new_file.txt\tProposed/,
+ );
+ });
+
+ it('should request confirmation with diff for an existing file (with corrected content)', async () => {
+ const filePath = path.join(rootDir, 'confirm_existing_file.txt');
+ const originalContent = 'Original content for confirmation.';
+ const proposedContent = 'Proposed replacement for confirmation.';
+ const correctedProposedContent =
+ 'Corrected replacement for confirmation.';
+ fs.writeFileSync(filePath, originalContent, 'utf8');
+
+ mockEnsureCorrectEdit.mockResolvedValue({
+ params: {
+ file_path: filePath,
+ old_string: originalContent,
+ new_string: correctedProposedContent,
+ },
+ occurrences: 1,
+ });
+
+ const params = { file_path: filePath, content: proposedContent };
+ const confirmation = (await tool.shouldConfirmExecute(
+ params,
+ abortSignal,
+ )) as ToolEditConfirmationDetails;
+
+ expect(mockEnsureCorrectEdit).toHaveBeenCalledWith(
+ originalContent,
+ {
+ old_string: originalContent,
+ new_string: proposedContent,
+ file_path: filePath,
+ },
+ mockGeminiClientInstance,
+ abortSignal,
+ );
+ expect(confirmation).toEqual(
+ expect.objectContaining({
+ title: `Confirm Write: ${path.basename(filePath)}`,
+ fileName: 'confirm_existing_file.txt',
+ fileDiff: expect.stringContaining(correctedProposedContent),
+ }),
+ );
+ expect(confirmation.fileDiff).toMatch(
+ originalContent.replace(/[.*+?^${}()|[\\]\\]/g, '\\$&'),
+ );
+ });
+ });
+
+ describe('execute', () => {
+ const abortSignal = new AbortController().signal;
+ it('should return error if params are invalid (relative path)', async () => {
+ const params = { file_path: 'relative.txt', content: 'test' };
+ const result = await tool.execute(params, abortSignal);
+ expect(result.llmContent).toMatch(/Error: Invalid parameters provided/);
+ expect(result.returnDisplay).toMatch(/Error: File path must be absolute/);
+ });
+
+ it('should return error if params are invalid (path outside root)', async () => {
+ const outsidePath = path.resolve(tempDir, 'outside-root.txt');
+ const params = { file_path: outsidePath, content: 'test' };
+ const result = await tool.execute(params, abortSignal);
+ expect(result.llmContent).toMatch(/Error: Invalid parameters provided/);
+ expect(result.returnDisplay).toMatch(
+ /Error: File path must be within the root directory/,
+ );
+ });
+
+ it('should return error if _getCorrectedFileContent returns an error during execute', async () => {
+ const filePath = path.join(rootDir, 'execute_error_file.txt');
+ const params = { file_path: filePath, content: 'test content' };
+ fs.writeFileSync(filePath, 'original', { mode: 0o000 });
+
+ const readError = new Error('Simulated read error for execute');
+ const originalReadFileSync = fs.readFileSync;
+ vi.spyOn(fs, 'readFileSync').mockImplementationOnce(() => {
+ throw readError;
+ });
+
+ const result = await tool.execute(params, abortSignal);
+ expect(result.llmContent).toMatch(/Error checking existing file/);
+ expect(result.returnDisplay).toMatch(
+ /Error checking existing file: Simulated read error for execute/,
+ );
+
+ vi.spyOn(fs, 'readFileSync').mockImplementation(originalReadFileSync);
+ fs.chmodSync(filePath, 0o600);
+ });
+
+ it('should write a new file with corrected content and return diff', async () => {
+ const filePath = path.join(rootDir, 'execute_new_corrected_file.txt');
+ const proposedContent = 'Proposed new content for execute.';
+ const correctedContent = 'Corrected new content for execute.';
+ mockEnsureCorrectFileContent.mockResolvedValue(correctedContent);
+
+ const params = { file_path: filePath, content: proposedContent };
+
+ const confirmDetails = await tool.shouldConfirmExecute(
+ params,
+ abortSignal,
+ );
+ if (typeof confirmDetails === 'object' && confirmDetails.onConfirm) {
+ await confirmDetails.onConfirm(ToolConfirmationOutcome.ProceedOnce);
+ }
+
+ const result = await tool.execute(params, abortSignal);
+
+ expect(mockEnsureCorrectFileContent).toHaveBeenCalledWith(
+ proposedContent,
+ mockGeminiClientInstance,
+ abortSignal,
+ );
+ expect(result.llmContent).toMatch(
+ /Successfully created and wrote to new file/,
+ );
+ expect(fs.existsSync(filePath)).toBe(true);
+ expect(fs.readFileSync(filePath, 'utf8')).toBe(correctedContent);
+ const display = result.returnDisplay as FileDiff;
+ expect(display.fileName).toBe('execute_new_corrected_file.txt');
+ expect(display.fileDiff).toMatch(
+ /--- execute_new_corrected_file.txt\tOriginal/,
+ );
+ expect(display.fileDiff).toMatch(
+ /\+\+\+ execute_new_corrected_file.txt\tWritten/,
+ );
+ expect(display.fileDiff).toMatch(
+ correctedContent.replace(/[.*+?^${}()|[\\]\\]/g, '\\$&'),
+ );
+ });
+
+ it('should overwrite an existing file with corrected content and return diff', async () => {
+ const filePath = path.join(
+ rootDir,
+ 'execute_existing_corrected_file.txt',
+ );
+ const initialContent = 'Initial content for execute.';
+ const proposedContent = 'Proposed overwrite for execute.';
+ const correctedProposedContent = 'Corrected overwrite for execute.';
+ fs.writeFileSync(filePath, initialContent, 'utf8');
+
+ mockEnsureCorrectEdit.mockResolvedValue({
+ params: {
+ file_path: filePath,
+ old_string: initialContent,
+ new_string: correctedProposedContent,
+ },
+ occurrences: 1,
+ });
+
+ const params = { file_path: filePath, content: proposedContent };
+
+ const confirmDetails = await tool.shouldConfirmExecute(
+ params,
+ abortSignal,
+ );
+ if (typeof confirmDetails === 'object' && confirmDetails.onConfirm) {
+ await confirmDetails.onConfirm(ToolConfirmationOutcome.ProceedOnce);
+ }
+
+ const result = await tool.execute(params, abortSignal);
+
+ expect(mockEnsureCorrectEdit).toHaveBeenCalledWith(
+ initialContent,
+ {
+ old_string: initialContent,
+ new_string: proposedContent,
+ file_path: filePath,
+ },
+ mockGeminiClientInstance,
+ abortSignal,
+ );
+ expect(result.llmContent).toMatch(/Successfully overwrote file/);
+ expect(fs.readFileSync(filePath, 'utf8')).toBe(correctedProposedContent);
+ const display = result.returnDisplay as FileDiff;
+ expect(display.fileName).toBe('execute_existing_corrected_file.txt');
+ expect(display.fileDiff).toMatch(
+ initialContent.replace(/[.*+?^${}()|[\\]\\]/g, '\\$&'),
+ );
+ expect(display.fileDiff).toMatch(
+ correctedProposedContent.replace(/[.*+?^${}()|[\\]\\]/g, '\\$&'),
+ );
+ });
+
+ it('should create directory if it does not exist', async () => {
+ const dirPath = path.join(rootDir, 'new_dir_for_write');
+ const filePath = path.join(dirPath, 'file_in_new_dir.txt');
+ const content = 'Content in new directory';
+ mockEnsureCorrectFileContent.mockResolvedValue(content); // Ensure this mock is active
+
+ const params = { file_path: filePath, content };
+ // Simulate confirmation if your logic requires it before execute, or remove if not needed for this path
+ const confirmDetails = await tool.shouldConfirmExecute(
+ params,
+ abortSignal,
+ );
+ if (typeof confirmDetails === 'object' && confirmDetails.onConfirm) {
+ await confirmDetails.onConfirm(ToolConfirmationOutcome.ProceedOnce);
+ }
+
+ await tool.execute(params, abortSignal);
+
+ expect(fs.existsSync(dirPath)).toBe(true);
+ expect(fs.statSync(dirPath).isDirectory()).toBe(true);
+ expect(fs.existsSync(filePath)).toBe(true);
+ expect(fs.readFileSync(filePath, 'utf8')).toBe(content);
+ });
+ });
+});