Files
2026-07-13 12:49:17 +08:00

132 lines
4.7 KiB
Python

# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates
# SPDX-License-Identifier: MIT
import unittest
from pathlib import Path
from unittest.mock import AsyncMock, patch
from trae_agent.tools.base import ToolCallArguments
from trae_agent.tools.edit_tool import TextEditorTool
class TestTextEditorTool(unittest.IsolatedAsyncioTestCase):
def setUp(self):
self.tool = TextEditorTool()
# Use current working directory for test paths
self.test_dir = Path.cwd() / "test_dir"
self.test_file = self.test_dir / "test_file.txt"
def mock_file_system(self, exists=True, is_dir=False, content=""):
"""Helper to mock file system operations"""
patcher = patch("pathlib.Path.exists", return_value=exists)
self.mock_exists = patcher.start()
self.addCleanup(patcher.stop)
patcher = patch("pathlib.Path.is_dir", return_value=is_dir)
self.mock_is_dir = patcher.start()
self.addCleanup(patcher.stop)
patcher = patch("pathlib.Path.read_text", return_value=content)
self.mock_read = patcher.start()
self.addCleanup(patcher.stop)
patcher = patch("pathlib.Path.write_text")
self.mock_write = patcher.start()
self.addCleanup(patcher.stop)
async def test_create_file(self):
self.mock_file_system(exists=False)
result = await self.tool.execute(
ToolCallArguments(
{
"command": "create",
"path": str(self.test_file),
"file_text": "new content",
}
)
)
self.mock_write.assert_called_once_with("new content")
self.assertIn("created successfully", result.output)
async def test_insert_line(self):
self.mock_file_system(content="line1\nline3")
result = await self.tool.execute(
ToolCallArguments(
{
"command": "insert",
"path": str(self.test_file),
"insert_line": 1,
"new_str": "line2",
}
)
)
self.mock_write.assert_called_once()
self.assertIn("edited", result.output)
async def test_invalid_command(self):
result = await self.tool.execute(
ToolCallArguments({"command": "invalid", "path": str(self.test_file.absolute())})
)
self.assertEqual(result.error_code, -1)
self.assertIn("Please provide a valid path", result.error)
async def test_str_replace_multiple_occurrences(self):
self.mock_file_system(content="dup\ndup\nline3")
result = await self.tool.execute(
ToolCallArguments(
{
"command": "str_replace",
"path": str(self.test_file),
"old_str": "dup",
"new_str": "new",
}
)
)
self.assertEqual(result.error_code, -1)
self.assertIn("Multiple occurrences", result.error or "")
async def test_str_replace_success(self):
self.mock_file_system(content="old_content\nline2")
result = await self.tool.execute(
ToolCallArguments(
{
"command": "str_replace",
"path": str(self.test_file),
"old_str": "old_content",
"new_str": "new_content",
}
)
)
self.mock_write.assert_called_once()
self.assertIn("edited", result.output)
async def test_view_directory(self):
self.mock_file_system(exists=True, is_dir=True)
with patch("trae_agent.tools.edit_tool.run", new_callable=AsyncMock) as mock_run:
mock_run.return_value = (0, "file1\nfile2", "")
result = await self.tool.execute(
ToolCallArguments({"command": "view", "path": str(self.test_dir)})
)
self.assertIn("files and directories", result.output)
async def test_view_file(self):
self.mock_file_system(exists=True, is_dir=False, content="line1\nline2\nline3")
result = await self.tool.execute(
ToolCallArguments({"command": "view", "path": str(self.test_file)})
)
self.assertRegex(result.output, r"\d+\s+line1")
async def test_relative_path(self):
result = await self.tool.execute(
ToolCallArguments({"command": "view", "path": "relative/path"})
)
self.assertIn("absolute path", result.error)
async def test_missing_parameters(self):
result = await self.tool.execute(ToolCallArguments({"command": "create"}))
self.assertIn("No path provided", result.error)
if __name__ == "__main__":
unittest.main()