diff options
Diffstat (limited to 'packages/core/src/tools/edit.test.ts')
| -rw-r--r-- | packages/core/src/tools/edit.test.ts | 499 |
1 files changed, 499 insertions, 0 deletions
diff --git a/packages/core/src/tools/edit.test.ts b/packages/core/src/tools/edit.test.ts new file mode 100644 index 00000000..08d0860d --- /dev/null +++ b/packages/core/src/tools/edit.test.ts @@ -0,0 +1,499 @@ +/** + * @license + * Copyright 2025 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +/* eslint-disable @typescript-eslint/no-explicit-any */ + +const mockEnsureCorrectEdit = vi.hoisted(() => vi.fn()); +const mockGenerateJson = vi.hoisted(() => vi.fn()); + +vi.mock('../utils/editCorrector.js', () => ({ + ensureCorrectEdit: mockEnsureCorrectEdit, +})); + +vi.mock('../core/client.js', () => ({ + GeminiClient: vi.fn().mockImplementation(() => ({ + generateJson: mockGenerateJson, + })), +})); + +import { describe, it, expect, beforeEach, afterEach, vi, Mock } from 'vitest'; +import { EditTool, EditToolParams } from './edit.js'; +import { FileDiff } from './tools.js'; +import path from 'path'; +import fs from 'fs'; +import os from 'os'; +import { Config } from '../config/config.js'; +import { Content, Part, SchemaUnion } from '@google/genai'; + +describe('EditTool', () => { + let tool: EditTool; + let tempDir: string; + let rootDir: string; + let mockConfig: Config; + + beforeEach(() => { + tempDir = fs.mkdtempSync(path.join(os.tmpdir(), 'edit-tool-test-')); + rootDir = path.join(tempDir, 'root'); + fs.mkdirSync(rootDir); + + mockConfig = { + getTargetDir: () => rootDir, + getAlwaysSkipModificationConfirmation: vi.fn(() => false), + setAlwaysSkipModificationConfirmation: vi.fn(), + // getGeminiConfig: () => ({ apiKey: 'test-api-key' }), // This was not a real Config method + // Add other properties/methods of Config if EditTool uses them + // Minimal other methods to satisfy Config type if needed by EditTool constructor or other direct uses: + getApiKey: () => 'test-api-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: () => ({}) as any, // Minimal mock for ToolRegistry + } as unknown as Config; + + // Reset mocks before each test + (mockConfig.getAlwaysSkipModificationConfirmation as Mock).mockClear(); + (mockConfig.setAlwaysSkipModificationConfirmation as Mock).mockClear(); + // Default to not skipping confirmation + (mockConfig.getAlwaysSkipModificationConfirmation as Mock).mockReturnValue( + false, + ); + + // Reset mocks and set default implementation for ensureCorrectEdit + mockEnsureCorrectEdit.mockReset(); + mockEnsureCorrectEdit.mockImplementation(async (currentContent, params) => { + let occurrences = 0; + if (params.old_string && currentContent) { + // Simple string counting for the mock + let index = currentContent.indexOf(params.old_string); + while (index !== -1) { + occurrences++; + index = currentContent.indexOf(params.old_string, index + 1); + } + } else if (params.old_string === '') { + occurrences = 0; // Creating a new file + } + return Promise.resolve({ params, occurrences }); + }); + + // Default mock for generateJson to return the snippet unchanged + mockGenerateJson.mockReset(); + mockGenerateJson.mockImplementation( + async (contents: Content[], schema: SchemaUnion) => { + // The problematic_snippet is the last part of the user's content + const userContent = contents.find((c: Content) => c.role === 'user'); + let promptText = ''; + if (userContent && userContent.parts) { + promptText = userContent.parts + .filter((p: Part) => typeof (p as any).text === 'string') + .map((p: Part) => (p as any).text) + .join('\n'); + } + const snippetMatch = promptText.match( + /Problematic target snippet:\n```\n([\s\S]*?)\n```/, + ); + const problematicSnippet = + snippetMatch && snippetMatch[1] ? snippetMatch[1] : ''; + + if (((schema as any).properties as any)?.corrected_target_snippet) { + return Promise.resolve({ + corrected_target_snippet: problematicSnippet, + }); + } + if (((schema as any).properties as any)?.corrected_new_string) { + // For new_string correction, we might need more sophisticated logic, + // but for now, returning original is a safe default if not specified by a test. + const originalNewStringMatch = promptText.match( + /original_new_string \(what was intended to replace original_old_string\):\n```\n([\s\S]*?)\n```/, + ); + const originalNewString = + originalNewStringMatch && originalNewStringMatch[1] + ? originalNewStringMatch[1] + : ''; + return Promise.resolve({ corrected_new_string: originalNewString }); + } + return Promise.resolve({}); // Default empty object if schema doesn't match + }, + ); + + tool = new EditTool(mockConfig); + }); + + afterEach(() => { + fs.rmSync(tempDir, { recursive: true, force: true }); + }); + + describe('_applyReplacement', () => { + // Access private method for testing + // Note: `tool` is initialized in `beforeEach` of the parent describe block + it('should return newString if isNewFile is true', () => { + expect((tool as any)._applyReplacement(null, 'old', 'new', true)).toBe( + 'new', + ); + expect( + (tool as any)._applyReplacement('existing', 'old', 'new', true), + ).toBe('new'); + }); + + it('should return newString if currentContent is null and oldString is empty (defensive)', () => { + expect((tool as any)._applyReplacement(null, '', 'new', false)).toBe( + 'new', + ); + }); + + it('should return empty string if currentContent is null and oldString is not empty (defensive)', () => { + expect((tool as any)._applyReplacement(null, 'old', 'new', false)).toBe( + '', + ); + }); + + it('should replace oldString with newString in currentContent', () => { + expect( + (tool as any)._applyReplacement( + 'hello old world old', + 'old', + 'new', + false, + ), + ).toBe('hello new world new'); + }); + + it('should return currentContent if oldString is empty and not a new file', () => { + expect( + (tool as any)._applyReplacement('hello world', '', 'new', false), + ).toBe('hello world'); + }); + }); + + describe('validateParams', () => { + it('should return null for valid params', () => { + const params: EditToolParams = { + file_path: path.join(rootDir, 'test.txt'), + old_string: 'old', + new_string: 'new', + }; + expect(tool.validateParams(params)).toBeNull(); + }); + + it('should return error for relative path', () => { + const params: EditToolParams = { + file_path: 'test.txt', + old_string: 'old', + new_string: 'new', + }; + expect(tool.validateParams(params)).toMatch(/File path must be absolute/); + }); + + it('should return error for path outside root', () => { + const params: EditToolParams = { + file_path: path.join(tempDir, 'outside-root.txt'), + old_string: 'old', + new_string: 'new', + }; + expect(tool.validateParams(params)).toMatch( + /File path must be within the root directory/, + ); + }); + }); + + describe('shouldConfirmExecute', () => { + const testFile = 'edit_me.txt'; + let filePath: string; + + beforeEach(() => { + filePath = path.join(rootDir, testFile); + }); + + it('should return false if params are invalid', async () => { + const params: EditToolParams = { + file_path: 'relative.txt', + old_string: 'old', + new_string: 'new', + }; + expect( + await tool.shouldConfirmExecute(params, new AbortController().signal), + ).toBe(false); + }); + + it('should request confirmation for valid edit', async () => { + fs.writeFileSync(filePath, 'some old content here'); + const params: EditToolParams = { + file_path: filePath, + old_string: 'old', + new_string: 'new', + }; + // ensureCorrectEdit will be called by shouldConfirmExecute + mockEnsureCorrectEdit.mockResolvedValueOnce({ params, occurrences: 1 }); + const confirmation = await tool.shouldConfirmExecute( + params, + new AbortController().signal, + ); + expect(confirmation).toEqual( + expect.objectContaining({ + title: `Confirm Edit: ${testFile}`, + fileName: testFile, + fileDiff: expect.any(String), + }), + ); + }); + + it('should return false if old_string is not found (ensureCorrectEdit returns 0)', async () => { + fs.writeFileSync(filePath, 'some content here'); + const params: EditToolParams = { + file_path: filePath, + old_string: 'not_found', + new_string: 'new', + }; + mockEnsureCorrectEdit.mockResolvedValueOnce({ params, occurrences: 0 }); + expect( + await tool.shouldConfirmExecute(params, new AbortController().signal), + ).toBe(false); + }); + + it('should return false if multiple occurrences of old_string are found (ensureCorrectEdit returns > 1)', async () => { + fs.writeFileSync(filePath, 'old old content here'); + const params: EditToolParams = { + file_path: filePath, + old_string: 'old', + new_string: 'new', + }; + mockEnsureCorrectEdit.mockResolvedValueOnce({ params, occurrences: 2 }); + expect( + await tool.shouldConfirmExecute(params, new AbortController().signal), + ).toBe(false); + }); + + it('should request confirmation for creating a new file (empty old_string)', async () => { + const newFileName = 'new_file.txt'; + const newFilePath = path.join(rootDir, newFileName); + const params: EditToolParams = { + file_path: newFilePath, + old_string: '', + new_string: 'new file content', + }; + // ensureCorrectEdit might not be called if old_string is empty, + // as shouldConfirmExecute handles this for diff generation. + // If it is called, it should return 0 occurrences for a new file. + mockEnsureCorrectEdit.mockResolvedValueOnce({ params, occurrences: 0 }); + const confirmation = await tool.shouldConfirmExecute( + params, + new AbortController().signal, + ); + expect(confirmation).toEqual( + expect.objectContaining({ + title: `Confirm Edit: ${newFileName}`, + fileName: newFileName, + fileDiff: expect.any(String), + }), + ); + }); + + it('should use corrected params from ensureCorrectEdit for diff generation', async () => { + const originalContent = 'This is the original string to be replaced.'; + const originalOldString = 'original string'; + const originalNewString = 'new string'; + + const correctedOldString = 'original string to be replaced'; // More specific + const correctedNewString = 'completely new string'; // Different replacement + const expectedFinalContent = 'This is the completely new string.'; + + fs.writeFileSync(filePath, originalContent); + const params: EditToolParams = { + file_path: filePath, + old_string: originalOldString, + new_string: originalNewString, + }; + + // The main beforeEach already calls mockEnsureCorrectEdit.mockReset() + // Set a specific mock for this test case + let mockCalled = false; + mockEnsureCorrectEdit.mockImplementationOnce( + async (content, p, client) => { + console.log('mockEnsureCorrectEdit CALLED IN TEST'); + mockCalled = true; + expect(content).toBe(originalContent); + expect(p).toBe(params); + expect(client).toBe((tool as any).client); + return { + params: { + file_path: filePath, + old_string: correctedOldString, + new_string: correctedNewString, + }, + occurrences: 1, + }; + }, + ); + + const confirmation = (await tool.shouldConfirmExecute( + params, + new AbortController().signal, + )) as FileDiff; + + expect(mockCalled).toBe(true); // Check if the mock implementation was run + // expect(mockEnsureCorrectEdit).toHaveBeenCalledWith(originalContent, params, expect.anything()); // Keep this commented for now + expect(confirmation).toEqual( + expect.objectContaining({ + title: `Confirm Edit: ${testFile}`, + fileName: testFile, + }), + ); + // Check that the diff is based on the corrected strings leading to the new state + expect(confirmation.fileDiff).toContain(`-${originalContent}`); + expect(confirmation.fileDiff).toContain(`+${expectedFinalContent}`); + + // Verify that applying the correctedOldString and correctedNewString to originalContent + // indeed produces the expectedFinalContent, which is what the diff should reflect. + const patchedContent = originalContent.replace( + correctedOldString, // This was the string identified by ensureCorrectEdit for replacement + correctedNewString, // This was the string identified by ensureCorrectEdit as the replacement + ); + expect(patchedContent).toBe(expectedFinalContent); + }); + }); + + describe('execute', () => { + const testFile = 'execute_me.txt'; + let filePath: string; + + beforeEach(() => { + filePath = path.join(rootDir, testFile); + // Default for execute tests, can be overridden + mockEnsureCorrectEdit.mockImplementation(async (content, params) => { + let occurrences = 0; + if (params.old_string && content) { + let index = content.indexOf(params.old_string); + while (index !== -1) { + occurrences++; + index = content.indexOf(params.old_string, index + 1); + } + } else if (params.old_string === '') { + occurrences = 0; + } + return { params, occurrences }; + }); + }); + + it('should return error if params are invalid', async () => { + const params: EditToolParams = { + file_path: 'relative.txt', + old_string: 'old', + new_string: 'new', + }; + const result = await tool.execute(params, new AbortController().signal); + expect(result.llmContent).toMatch(/Error: Invalid parameters provided/); + expect(result.returnDisplay).toMatch(/Error: File path must be absolute/); + }); + + it('should edit an existing file and return diff with fileName', async () => { + const initialContent = 'This is some old text.'; + const newContent = 'This is some new text.'; // old -> new + fs.writeFileSync(filePath, initialContent, 'utf8'); + const params: EditToolParams = { + file_path: filePath, + old_string: 'old', + new_string: 'new', + }; + + // Specific mock for this test's execution path in calculateEdit + // ensureCorrectEdit is NOT called by calculateEdit, only by shouldConfirmExecute + // So, the default mockEnsureCorrectEdit should correctly return 1 occurrence for 'old' in initialContent + + // Simulate confirmation by setting shouldAlwaysEdit + (tool as any).shouldAlwaysEdit = true; + + const result = await tool.execute(params, new AbortController().signal); + + (tool as any).shouldAlwaysEdit = false; // Reset for other tests + + expect(result.llmContent).toMatch(/Successfully modified file/); + expect(fs.readFileSync(filePath, 'utf8')).toBe(newContent); + const display = result.returnDisplay as FileDiff; + expect(display.fileDiff).toMatch(initialContent); + expect(display.fileDiff).toMatch(newContent); + expect(display.fileName).toBe(testFile); + }); + + it('should create a new file if old_string is empty and file does not exist, and return created message', async () => { + const newFileName = 'brand_new_file.txt'; + const newFilePath = path.join(rootDir, newFileName); + const fileContent = 'Content for the new file.'; + const params: EditToolParams = { + file_path: newFilePath, + old_string: '', + new_string: fileContent, + }; + + ( + mockConfig.getAlwaysSkipModificationConfirmation as Mock + ).mockReturnValueOnce(true); + const result = await tool.execute(params, new AbortController().signal); + + expect(result.llmContent).toMatch(/Created new file/); + expect(fs.existsSync(newFilePath)).toBe(true); + expect(fs.readFileSync(newFilePath, 'utf8')).toBe(fileContent); + expect(result.returnDisplay).toBe(`Created ${newFileName}`); + }); + + it('should return error if old_string is not found in file', async () => { + fs.writeFileSync(filePath, 'Some content.', 'utf8'); + const params: EditToolParams = { + file_path: filePath, + old_string: 'nonexistent', + new_string: 'replacement', + }; + // The default mockEnsureCorrectEdit will return 0 occurrences for 'nonexistent' + const result = await tool.execute(params, new AbortController().signal); + expect(result.llmContent).toMatch( + /0 occurrences found for old_string in/, + ); + expect(result.returnDisplay).toMatch( + /Failed to edit, could not find the string to replace./, + ); + }); + + it('should return error if multiple occurrences of old_string are found', async () => { + fs.writeFileSync(filePath, 'multiple old old strings', 'utf8'); + const params: EditToolParams = { + file_path: filePath, + old_string: 'old', + new_string: 'new', + }; + // The default mockEnsureCorrectEdit will return 2 occurrences for 'old' + const result = await tool.execute(params, new AbortController().signal); + expect(result.llmContent).toMatch( + /Expected 1 occurrences but found 2 for old_string in file/, + ); + expect(result.returnDisplay).toMatch( + /Failed to edit, expected 1 occurrence\(s\) but found 2/, + ); + }); + + it('should return error if trying to create a file that already exists (empty old_string)', async () => { + fs.writeFileSync(filePath, 'Existing content', 'utf8'); + const params: EditToolParams = { + file_path: filePath, + old_string: '', + new_string: 'new content', + }; + const result = await tool.execute(params, new AbortController().signal); + expect(result.llmContent).toMatch(/File already exists, cannot create/); + expect(result.returnDisplay).toMatch( + /Attempted to create a file that already exists/, + ); + }); + }); +}); |