From 3281cbc835ae77d0c33662485505651d01a984cc Mon Sep 17 00:00:00 2001 From: Taylor Mullen Date: Sun, 25 May 2025 14:59:45 -0700 Subject: [PATCH] Fix(test): Improve write-file and editCorrector test suites - Enhanced by: - Mocking and utilities (, ) to allow for more focused unit testing of . - Adding comprehensive tests for the private method, covering new and existing file scenarios, as well as error handling. - Expanding tests for and to verify behavior with the new content correction logic, including diff generation and directory creation. - Refined by: - Introducing robust mocking for and its methods (, , ) to simulate LLM interactions accurately. - Adding extensive test scenarios for , categorized by how matches and how is processed, including direct matches, unescaping, and LLM-based corrections. - Including tests for edge cases like no matches or multiple matches. - Adding a utility for better test isolation. Final fix for https://github.com/google-gemini/gemini-cli/issues/484 --- packages/server/src/tools/write-file.test.ts | 397 +++++++++++++++-- .../server/src/utils/editCorrector.test.ts | 415 ++++++++++++++++-- 2 files changed, 742 insertions(+), 70 deletions(-) diff --git a/packages/server/src/tools/write-file.test.ts b/packages/server/src/tools/write-file.test.ts index 9a690521..83e75a33 100644 --- a/packages/server/src/tools/write-file.test.ts +++ b/packages/server/src/tools/write-file.test.ts @@ -4,17 +4,50 @@ * SPDX-License-Identifier: Apache-2.0 */ -import { describe, it, expect, beforeEach, afterEach, vi } from 'vitest'; +import { + describe, + it, + expect, + beforeEach, + afterEach, + vi, + type Mocked, +} from 'vitest'; import { WriteFileTool } from './write-file.js'; -import { FileDiff, ToolConfirmationOutcome } from './tools.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'; // Added import +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; +const mockEnsureCorrectEdit = vi.fn(); +const mockEnsureCorrectFileContent = vi.fn(); + +// 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, @@ -42,13 +75,14 @@ const mockConfigInternal = { }) 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 (for testing boundary conditions) + // Create a unique temporary directory for files created outside the root tempDir = fs.mkdtempSync( path.join(os.tmpdir(), 'write-file-test-external-'), ); @@ -56,12 +90,39 @@ describe('WriteFileTool', () => { if (!fs.existsSync(rootDir)) { fs.mkdirSync(rootDir, { recursive: true }); } + + // Setup GeminiClient mock + mockGeminiClientInstance = new (vi.mocked(GeminiClient))( + mockConfig, + ) as Mocked; + vi.mocked(GeminiClient).mockImplementation(() => mockGeminiClientInstance); + tool = new WriteFileTool(mockConfig); - // Reset mocks before each test that might use them for confirmation logic + + // 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, + ): Promise => + Promise.resolve({ + params: { ...params, new_string: params.new_string ?? '' }, + occurrences: 1, + }), + ); + mockEnsureCorrectFileContent.mockImplementation( + async (content: string, _client: GeminiClient): Promise => + Promise.resolve(content ?? ''), + ); }); afterEach(() => { @@ -72,6 +133,7 @@ describe('WriteFileTool', () => { if (fs.existsSync(rootDir)) { fs.rmSync(rootDir, { recursive: true, force: true }); } + vi.clearAllMocks(); }); describe('validateToolParams', () => { @@ -101,27 +163,113 @@ describe('WriteFileTool', () => { ); }); - it('should return error for path that is the root itself', () => { + it('should return error if path is a directory', () => { + const dirAsFilePath = path.join(rootDir, 'a_directory'); + fs.mkdirSync(dirAsFilePath); const params = { - file_path: rootDir, // Attempting to write to the root directory itself + file_path: dirAsFilePath, content: 'hello', }; - // With the new validation, this should now return an error as rootDir is a directory. expect(tool.validateToolParams(params)).toMatch( - `Path is a directory, not a file: ${rootDir}`, + `Path is a directory, not a file: ${dirAsFilePath}`, ); }); + }); - it('should return error for path that is just / and root is not /', () => { - const params = { file_path: path.resolve('/'), content: 'hello' }; - if (rootDir === path.resolve('/')) { - // This case would only occur if the test runner somehow sets rootDir to actual '/', which is highly unlikely and unsafe. - expect(tool.validateToolParams(params)).toBeNull(); - } else { - expect(tool.validateToolParams(params)).toMatch( - /File path must be within the root directory/, - ); - } + 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.'; + // 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, + ); + + expect(mockEnsureCorrectFileContent).toHaveBeenCalledWith( + proposedContent, + mockGeminiClientInstance, + ); + 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.'; + 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, + ); + + expect(mockEnsureCorrectEdit).toHaveBeenCalledWith( + originalContent, + { + old_string: originalContent, + new_string: proposedContent, + file_path: filePath, + }, + mockGeminiClientInstance, + ); + 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'; + 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, + ); + + 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); }); }); @@ -139,17 +287,96 @@ describe('WriteFileTool', () => { expect(confirmation).toBe(false); }); - it('should request confirmation for valid params if file does not exist', async () => { - const filePath = path.join(rootDir, 'new_file.txt'); - const params = { file_path: filePath, content: 'new content' }; + 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); + 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, + )) as ToolEditConfirmationDetails; + + expect(mockEnsureCorrectFileContent).toHaveBeenCalledWith( + proposedContent, + mockGeminiClientInstance, + ); expect(confirmation).toEqual( expect.objectContaining({ title: `Confirm Write: ${path.basename(filePath)}`, - fileName: 'new_file.txt', - fileDiff: expect.any(String), + 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'); + + // Ensure this mock is active and returns the correct structure + 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, + )) as ToolEditConfirmationDetails; + + expect(mockEnsureCorrectEdit).toHaveBeenCalledWith( + originalContent, + { + old_string: originalContent, + new_string: proposedContent, + file_path: filePath, + }, + mockGeminiClientInstance, + ); + 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, '\\$&'), + ); }); }); @@ -171,10 +398,34 @@ describe('WriteFileTool', () => { ); }); - it('should write a new file and return diff', async () => { - const filePath = path.join(rootDir, 'execute_new_file.txt'); - const content = 'Hello from execute!'; - const params = { file_path: filePath, content }; + 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, new AbortController().signal); + 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); if (typeof confirmDetails === 'object' && confirmDetails.onConfirm) { @@ -183,26 +434,48 @@ describe('WriteFileTool', () => { const result = await tool.execute(params, new AbortController().signal); + expect(mockEnsureCorrectFileContent).toHaveBeenCalledWith( + proposedContent, + mockGeminiClientInstance, + ); expect(result.llmContent).toMatch( /Successfully created and wrote to new file/, ); expect(fs.existsSync(filePath)).toBe(true); - expect(fs.readFileSync(filePath, 'utf8')).toBe(content); - const display = result.returnDisplay as FileDiff; // Type assertion - expect(display.fileName).toBe('execute_new_file.txt'); - // For new files, the diff will include the filename in the "Original" header - expect(display.fileDiff).toMatch(/--- execute_new_file.txt\tOriginal/); - expect(display.fileDiff).toMatch(/\+\+\+ execute_new_file.txt\tWritten/); - expect(display.fileDiff).toMatch(content); + 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 and return diff', async () => { - const filePath = path.join(rootDir, 'execute_existing_file.txt'); - const initialContent = 'Initial content.'; - const newContent = 'Overwritten content!'; + 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'); - const params = { file_path: filePath, content: newContent }; + 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); if (typeof confirmDetails === 'object' && confirmDetails.onConfirm) { @@ -211,12 +484,46 @@ describe('WriteFileTool', () => { const result = await tool.execute(params, new AbortController().signal); + expect(mockEnsureCorrectEdit).toHaveBeenCalledWith( + initialContent, + { + old_string: initialContent, + new_string: proposedContent, + file_path: filePath, + }, + mockGeminiClientInstance, + ); expect(result.llmContent).toMatch(/Successfully overwrote file/); - expect(fs.readFileSync(filePath, 'utf8')).toBe(newContent); - const display = result.returnDisplay as FileDiff; // Type assertion - expect(display.fileName).toBe('execute_existing_file.txt'); - expect(display.fileDiff).toMatch(initialContent); - expect(display.fileDiff).toMatch(newContent); + 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); + if (typeof confirmDetails === 'object' && confirmDetails.onConfirm) { + await confirmDetails.onConfirm(ToolConfirmationOutcome.ProceedOnce); + } + + await tool.execute(params, new AbortController().signal); + + 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); }); }); }); diff --git a/packages/server/src/utils/editCorrector.test.ts b/packages/server/src/utils/editCorrector.test.ts index 8b27bdf1..27c9ffe8 100644 --- a/packages/server/src/utils/editCorrector.test.ts +++ b/packages/server/src/utils/editCorrector.test.ts @@ -4,44 +4,68 @@ * SPDX-License-Identifier: Apache-2.0 */ -import { describe, it, expect } from 'vitest'; +/* eslint-disable @typescript-eslint/no-explicit-any */ +import { vi, describe, it, expect, beforeEach, type Mocked } from 'vitest'; + +// MOCKS +let callCount = 0; +const mockResponses: any[] = []; + +let mockGenerateJson: any; +let mockStartChat: any; +let mockSendMessageStream: any; + +vi.mock('../core/client.js', () => ({ + GeminiClient: vi.fn().mockImplementation(function ( + this: any, + _config: Config, + ) { + this.generateJson = (...params: any[]) => mockGenerateJson(...params); // Corrected: use mockGenerateJson + this.startChat = (...params: any[]) => mockStartChat(...params); // Corrected: use mockStartChat + this.sendMessageStream = (...params: any[]) => + mockSendMessageStream(...params); // Corrected: use mockSendMessageStream + return this; + }), +})); +// END MOCKS + import { countOccurrences, + ensureCorrectEdit, unescapeStringForGeminiBug, + resetEditCorrectorCaches_TEST_ONLY, } from './editCorrector.js'; +import { GeminiClient } from '../core/client.js'; +import type { Config } from '../config/config.js'; +import { ToolRegistry } from '../tools/tool-registry.js'; + +vi.mock('../tools/tool-registry.js'); describe('editCorrector', () => { describe('countOccurrences', () => { it('should return 0 for empty string', () => { expect(countOccurrences('', 'a')).toBe(0); }); - it('should return 0 for empty substring', () => { expect(countOccurrences('abc', '')).toBe(0); }); - it('should return 0 if substring is not found', () => { expect(countOccurrences('abc', 'd')).toBe(0); }); - it('should return 1 if substring is found once', () => { expect(countOccurrences('abc', 'b')).toBe(1); }); - it('should return correct count for multiple occurrences', () => { expect(countOccurrences('ababa', 'a')).toBe(3); expect(countOccurrences('ababab', 'ab')).toBe(3); }); - it('should count non-overlapping occurrences', () => { - expect(countOccurrences('aaaaa', 'aa')).toBe(2); // Non-overlapping: aa_aa_ - expect(countOccurrences('ababab', 'aba')).toBe(1); // Non-overlapping: aba_ab -> 1 + expect(countOccurrences('aaaaa', 'aa')).toBe(2); + expect(countOccurrences('ababab', 'aba')).toBe(1); }); - it('should correctly count occurrences when substring is longer', () => { expect(countOccurrences('abc', 'abcdef')).toBe(0); }); - it('should be case sensitive', () => { expect(countOccurrences('abcABC', 'a')).toBe(1); expect(countOccurrences('abcABC', 'A')).toBe(1); @@ -56,62 +80,403 @@ describe('editCorrector', () => { expect(unescapeStringForGeminiBug('\\"')).toBe('"'); expect(unescapeStringForGeminiBug('\\`')).toBe('`'); }); - it('should handle multiple escaped sequences', () => { expect(unescapeStringForGeminiBug('Hello\\nWorld\\tTest')).toBe( 'Hello\nWorld\tTest', ); }); - it('should not alter already correct sequences', () => { expect(unescapeStringForGeminiBug('\n')).toBe('\n'); expect(unescapeStringForGeminiBug('Correct string')).toBe( 'Correct string', ); }); - it('should handle mixed correct and incorrect sequences', () => { expect(unescapeStringForGeminiBug('\\nCorrect\t\\`')).toBe( '\nCorrect\t`', ); }); - it('should handle backslash followed by actual newline character', () => { expect(unescapeStringForGeminiBug('\\\n')).toBe('\n'); expect(unescapeStringForGeminiBug('First line\\\nSecond line')).toBe( 'First line\nSecond line', ); }); - it('should handle multiple backslashes before an escapable character', () => { - expect(unescapeStringForGeminiBug('\\\\n')).toBe('\n'); // \\n -> \n - expect(unescapeStringForGeminiBug('\\\\\\t')).toBe('\t'); // \\\t -> \t - expect(unescapeStringForGeminiBug('\\\\\\\\`')).toBe('`'); // \\\\` -> ` + expect(unescapeStringForGeminiBug('\\\\n')).toBe('\n'); + expect(unescapeStringForGeminiBug('\\\\\\t')).toBe('\t'); + expect(unescapeStringForGeminiBug('\\\\\\\\`')).toBe('`'); }); - it('should return empty string for empty input', () => { expect(unescapeStringForGeminiBug('')).toBe(''); }); - it('should not alter strings with no targeted escape sequences', () => { expect(unescapeStringForGeminiBug('abc def')).toBe('abc def'); - // \\F and \\S are not targeted escapes, so they should remain as \\F and \\S expect(unescapeStringForGeminiBug('C:\\Folder\\File')).toBe( 'C:\\Folder\\File', ); }); - it('should correctly process strings with some targeted escapes', () => { - // \\U is not targeted, \\n is. expect(unescapeStringForGeminiBug('C:\\Users\\name')).toBe( 'C:\\Users\name', ); }); - it('should handle complex cases with mixed slashes and characters', () => { expect( - unescapeStringForGeminiBug('\\\\\\nLine1\\\nLine2\\tTab\\\\`Tick\\"'), + unescapeStringForGeminiBug('\\\\\\\nLine1\\\nLine2\\tTab\\\\`Tick\\"'), ).toBe('\nLine1\nLine2\tTab`Tick"'); }); }); + + describe('ensureCorrectEdit', () => { + let mockGeminiClientInstance: Mocked; + let mockToolRegistry: Mocked; + let mockConfigInstance: Config; + + beforeEach(() => { + mockToolRegistry = new ToolRegistry({} as Config) as Mocked; + const configParams = { + apiKey: 'test-api-key', + model: 'test-model', + sandbox: false as boolean | string, + targetDir: '/test', + debugMode: false, + question: undefined as string | undefined, + fullContext: false, + coreTools: undefined as string[] | undefined, + toolDiscoveryCommand: undefined as string | undefined, + toolCallCommand: undefined as string | undefined, + mcpServerCommand: undefined as string | undefined, + mcpServers: undefined as Record | undefined, + userAgent: 'test-agent', + userMemory: '', + geminiMdFileCount: 0, + alwaysSkipModificationConfirmation: false, + }; + mockConfigInstance = { + ...configParams, + getApiKey: vi.fn(() => configParams.apiKey), + getModel: vi.fn(() => configParams.model), + getSandbox: vi.fn(() => configParams.sandbox), + getTargetDir: vi.fn(() => configParams.targetDir), + getToolRegistry: vi.fn(() => mockToolRegistry), + getDebugMode: vi.fn(() => configParams.debugMode), + getQuestion: vi.fn(() => configParams.question), + getFullContext: vi.fn(() => configParams.fullContext), + getCoreTools: vi.fn(() => configParams.coreTools), + getToolDiscoveryCommand: vi.fn(() => configParams.toolDiscoveryCommand), + getToolCallCommand: vi.fn(() => configParams.toolCallCommand), + getMcpServerCommand: vi.fn(() => configParams.mcpServerCommand), + getMcpServers: vi.fn(() => configParams.mcpServers), + getUserAgent: vi.fn(() => configParams.userAgent), + getUserMemory: vi.fn(() => configParams.userMemory), + setUserMemory: vi.fn((mem: string) => { + configParams.userMemory = mem; + }), + getGeminiMdFileCount: vi.fn(() => configParams.geminiMdFileCount), + setGeminiMdFileCount: vi.fn((count: number) => { + configParams.geminiMdFileCount = count; + }), + getAlwaysSkipModificationConfirmation: vi.fn( + () => configParams.alwaysSkipModificationConfirmation, + ), + setAlwaysSkipModificationConfirmation: vi.fn((skip: boolean) => { + configParams.alwaysSkipModificationConfirmation = skip; + }), + } as unknown as Config; + + callCount = 0; + mockResponses.length = 0; + mockGenerateJson = vi.fn().mockImplementation(() => { + const response = mockResponses[callCount]; + callCount++; + if (response === undefined) return Promise.resolve({}); + return Promise.resolve(response); + }); + mockStartChat = vi.fn(); + mockSendMessageStream = vi.fn(); + + mockGeminiClientInstance = new GeminiClient( + mockConfigInstance, + ) as Mocked; + resetEditCorrectorCaches_TEST_ONLY(); + }); + + describe('Scenario Group 1: originalParams.old_string matches currentContent directly', () => { + it('Test 1.1: old_string (no literal \\), new_string (escaped by Gemini) -> new_string unescaped', async () => { + const currentContent = 'This is a test string to find me.'; + const originalParams = { + file_path: '/test/file.txt', + old_string: 'find me', + new_string: 'replace with \\"this\\"', + }; + mockResponses.push({ + corrected_new_string_escaping: 'replace with "this"', + }); + const result = await ensureCorrectEdit( + currentContent, + originalParams, + mockGeminiClientInstance, + ); + expect(mockGenerateJson).toHaveBeenCalledTimes(1); + expect(result.params.new_string).toBe('replace with "this"'); + expect(result.params.old_string).toBe('find me'); + expect(result.occurrences).toBe(1); + }); + it('Test 1.2: old_string (no literal \\), new_string (correctly formatted) -> new_string unchanged', async () => { + const currentContent = 'This is a test string to find me.'; + const originalParams = { + file_path: '/test/file.txt', + old_string: 'find me', + new_string: 'replace with this', + }; + const result = await ensureCorrectEdit( + currentContent, + originalParams, + mockGeminiClientInstance, + ); + expect(mockGenerateJson).toHaveBeenCalledTimes(0); + expect(result.params.new_string).toBe('replace with this'); + expect(result.params.old_string).toBe('find me'); + expect(result.occurrences).toBe(1); + }); + it('Test 1.3: old_string (with literal \\), new_string (escaped by Gemini) -> new_string unchanged (still escaped)', async () => { + const currentContent = 'This is a test string to find\\me.'; + const originalParams = { + file_path: '/test/file.txt', + old_string: 'find\\me', + new_string: 'replace with \\"this\\"', + }; + mockResponses.push({ + corrected_new_string_escaping: 'replace with "this"', + }); + const result = await ensureCorrectEdit( + currentContent, + originalParams, + mockGeminiClientInstance, + ); + expect(mockGenerateJson).toHaveBeenCalledTimes(1); + expect(result.params.new_string).toBe('replace with "this"'); + expect(result.params.old_string).toBe('find\\me'); + expect(result.occurrences).toBe(1); + }); + it('Test 1.4: old_string (with literal \\), new_string (correctly formatted) -> new_string unchanged', async () => { + const currentContent = 'This is a test string to find\\me.'; + const originalParams = { + file_path: '/test/file.txt', + old_string: 'find\\me', + new_string: 'replace with this', + }; + const result = await ensureCorrectEdit( + currentContent, + originalParams, + mockGeminiClientInstance, + ); + expect(mockGenerateJson).toHaveBeenCalledTimes(0); + expect(result.params.new_string).toBe('replace with this'); + expect(result.params.old_string).toBe('find\\me'); + expect(result.occurrences).toBe(1); + }); + }); + + describe('Scenario Group 2: originalParams.old_string does NOT match, but unescapeStringForGeminiBug(originalParams.old_string) DOES match', () => { + it('Test 2.1: old_string (over-escaped, no intended literal \\), new_string (escaped by Gemini) -> new_string unescaped', async () => { + const currentContent = 'This is a test string to find "me".'; + const originalParams = { + file_path: '/test/file.txt', + old_string: 'find \\"me\\"', + new_string: 'replace with \\"this\\"', + }; + mockResponses.push({ corrected_new_string: 'replace with "this"' }); + const result = await ensureCorrectEdit( + currentContent, + originalParams, + mockGeminiClientInstance, + ); + expect(mockGenerateJson).toHaveBeenCalledTimes(1); + expect(result.params.new_string).toBe('replace with "this"'); + expect(result.params.old_string).toBe('find "me"'); + expect(result.occurrences).toBe(1); + }); + it('Test 2.2: old_string (over-escaped, no intended literal \\), new_string (correctly formatted) -> new_string unescaped (harmlessly)', async () => { + const currentContent = 'This is a test string to find "me".'; + const originalParams = { + file_path: '/test/file.txt', + old_string: 'find \\"me\\"', + new_string: 'replace with this', + }; + const result = await ensureCorrectEdit( + currentContent, + originalParams, + mockGeminiClientInstance, + ); + expect(mockGenerateJson).toHaveBeenCalledTimes(0); + expect(result.params.new_string).toBe('replace with this'); + expect(result.params.old_string).toBe('find "me"'); + expect(result.occurrences).toBe(1); + }); + it('Test 2.3: old_string (over-escaped, with intended literal \\), new_string (simple) -> new_string corrected', async () => { + const currentContent = 'This is a test string to find \\me.'; + const originalParams = { + file_path: '/test/file.txt', + old_string: 'find \\\\me', + new_string: 'replace with foobar', + }; + mockResponses.push({ + corrected_target_snippet: 'find \\me', + }); + const result = await ensureCorrectEdit( + currentContent, + originalParams, + mockGeminiClientInstance, + ); + expect(mockGenerateJson).toHaveBeenCalledTimes(1); + expect(result.params.new_string).toBe('replace with foobar'); + expect(result.params.old_string).toBe('find \\me'); + expect(result.occurrences).toBe(1); + }); + }); + + describe('Scenario Group 3: LLM Correction Path', () => { + it('Test 3.1: old_string (no literal \\), new_string (escaped by Gemini), LLM re-escapes new_string -> final new_string is double unescaped', async () => { + const currentContent = 'This is a test string to corrected find me.'; + const originalParams = { + file_path: '/test/file.txt', + old_string: 'find me', + new_string: 'replace with \\\\"this\\\\"', + }; + const llmNewString = 'LLM says replace with "that"'; + mockResponses.push({ corrected_new_string_escaping: llmNewString }); + const result = await ensureCorrectEdit( + currentContent, + originalParams, + mockGeminiClientInstance, + ); + expect(mockGenerateJson).toHaveBeenCalledTimes(1); + expect(result.params.new_string).toBe(llmNewString); + expect(result.params.old_string).toBe('find me'); + expect(result.occurrences).toBe(1); + }); + it('Test 3.2: old_string (with literal \\), new_string (escaped by Gemini), LLM re-escapes new_string -> final new_string is unescaped once', async () => { + const currentContent = 'This is a test string to corrected find me.'; + const originalParams = { + file_path: '/test/file.txt', + old_string: 'find\\me', + new_string: 'replace with \\\\"this\\\\"', + }; + const llmCorrectedOldString = 'corrected find me'; + const llmNewString = 'LLM says replace with "that"'; + mockResponses.push({ corrected_target_snippet: llmCorrectedOldString }); + mockResponses.push({ corrected_new_string: llmNewString }); + const result = await ensureCorrectEdit( + currentContent, + originalParams, + mockGeminiClientInstance, + ); + expect(mockGenerateJson).toHaveBeenCalledTimes(2); + expect(result.params.new_string).toBe(llmNewString); + expect(result.params.old_string).toBe(llmCorrectedOldString); + expect(result.occurrences).toBe(1); + }); + it('Test 3.3: old_string needs LLM, new_string is fine -> old_string corrected, new_string original', async () => { + const currentContent = 'This is a test string to be corrected.'; + const originalParams = { + file_path: '/test/file.txt', + old_string: 'fiiind me', + new_string: 'replace with "this"', + }; + const llmCorrectedOldString = 'to be corrected'; + mockResponses.push({ corrected_target_snippet: llmCorrectedOldString }); + const result = await ensureCorrectEdit( + currentContent, + originalParams, + mockGeminiClientInstance, + ); + expect(mockGenerateJson).toHaveBeenCalledTimes(1); + expect(result.params.new_string).toBe('replace with "this"'); + expect(result.params.old_string).toBe(llmCorrectedOldString); + expect(result.occurrences).toBe(1); + }); + it('Test 3.4: LLM correction path, correctNewString returns the originalNewString it was passed (which was unescaped) -> final new_string is unescaped', async () => { + const currentContent = 'This is a test string to corrected find me.'; + const originalParams = { + file_path: '/test/file.txt', + old_string: 'find me', + new_string: 'replace with \\\\"this\\\\"', + }; + const newStringForLLMAndReturnedByLLM = 'replace with "this"'; + mockResponses.push({ + corrected_new_string_escaping: newStringForLLMAndReturnedByLLM, + }); + const result = await ensureCorrectEdit( + currentContent, + originalParams, + mockGeminiClientInstance, + ); + expect(mockGenerateJson).toHaveBeenCalledTimes(1); + expect(result.params.new_string).toBe(newStringForLLMAndReturnedByLLM); + expect(result.occurrences).toBe(1); + }); + }); + + describe('Scenario Group 4: No Match Found / Multiple Matches', () => { + it('Test 4.1: No version of old_string (original, unescaped, LLM-corrected) matches -> returns original params, 0 occurrences', async () => { + const currentContent = 'This content has nothing to find.'; + const originalParams = { + file_path: '/test/file.txt', + old_string: 'nonexistent string', + new_string: 'some new string', + }; + mockResponses.push({ corrected_target_snippet: 'still nonexistent' }); + const result = await ensureCorrectEdit( + currentContent, + originalParams, + mockGeminiClientInstance, + ); + expect(mockGenerateJson).toHaveBeenCalledTimes(1); + expect(result.params).toEqual(originalParams); + expect(result.occurrences).toBe(0); + }); + it('Test 4.2: unescapedOldStringAttempt results in >1 occurrences -> returns original params, count occurrences', async () => { + const currentContent = + 'This content has find "me" and also find "me" again.'; + const originalParams = { + file_path: '/test/file.txt', + old_string: 'find "me"', + new_string: 'some new string', + }; + const result = await ensureCorrectEdit( + currentContent, + originalParams, + mockGeminiClientInstance, + ); + expect(mockGenerateJson).toHaveBeenCalledTimes(0); + expect(result.params).toEqual(originalParams); + expect(result.occurrences).toBe(2); + }); + }); + + describe('Scenario Group 5: Specific unescapeStringForGeminiBug checks (integrated into ensureCorrectEdit)', () => { + it('Test 5.1: old_string needs LLM to become currentContent, new_string also needs correction', async () => { + const currentContent = 'const x = "a\\nbc\\\\"def\\\\"'; + const originalParams = { + file_path: '/test/file.txt', + old_string: 'const x = \\\\"a\\\\nbc\\\\\\\\"def\\\\\\\\"', + new_string: 'const y = \\\\"new\\\\nval\\\\\\\\"content\\\\\\\\"', + }; + const expectedFinalNewString = 'const y = "new\\nval\\\\"content\\\\"'; + mockResponses.push({ corrected_target_snippet: currentContent }); + mockResponses.push({ corrected_new_string: expectedFinalNewString }); + const result = await ensureCorrectEdit( + currentContent, + originalParams, + mockGeminiClientInstance, + ); + expect(mockGenerateJson).toHaveBeenCalledTimes(2); + expect(result.params.old_string).toBe(currentContent); + expect(result.params.new_string).toBe(expectedFinalNewString); + expect(result.occurrences).toBe(1); + }); + }); + }); });