improve file editing and add notebook tool
Enhance file tools with read tracking, PDF support, safer path handling, smarter edit matching/diagnostics, and introduce notebook_edit with tests.
This commit is contained in:
Xubin Ren
@@ -0,0 +1,423 @@
|
||||
"""Tests for advanced EditFileTool enhancements inspired by claude-code:
|
||||
- Delete-line newline cleanup
|
||||
- Smart quote normalization (curly ↔ straight)
|
||||
- Quote style preservation in replacements
|
||||
- Indentation preservation when fallback match is trimmed
|
||||
- Trailing whitespace stripping for new_text
|
||||
- File size protection
|
||||
- Stale detection with content-equality fallback
|
||||
"""
|
||||
|
||||
import os
|
||||
import time
|
||||
|
||||
import pytest
|
||||
|
||||
from nanobot.agent.tools.filesystem import EditFileTool, ReadFileTool, _find_match
|
||||
from nanobot.agent.tools import file_state
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _clear_file_state():
|
||||
file_state.clear()
|
||||
yield
|
||||
file_state.clear()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Delete-line newline cleanup
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestDeleteLineCleanup:
|
||||
"""When new_text='' and deleting a line, trailing newline should be consumed."""
|
||||
|
||||
@pytest.fixture()
|
||||
def tool(self, tmp_path):
|
||||
return EditFileTool(workspace=tmp_path)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_delete_line_consumes_trailing_newline(self, tool, tmp_path):
|
||||
f = tmp_path / "a.py"
|
||||
f.write_text("line1\nline2\nline3\n", encoding="utf-8")
|
||||
result = await tool.execute(path=str(f), old_text="line2", new_text="")
|
||||
assert "Successfully" in result
|
||||
content = f.read_text()
|
||||
# Should not leave a blank line where line2 was
|
||||
assert content == "line1\nline3\n"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_delete_line_with_explicit_newline_in_old_text(self, tool, tmp_path):
|
||||
f = tmp_path / "a.py"
|
||||
f.write_text("line1\nline2\nline3\n", encoding="utf-8")
|
||||
result = await tool.execute(path=str(f), old_text="line2\n", new_text="")
|
||||
assert "Successfully" in result
|
||||
assert f.read_text() == "line1\nline3\n"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_delete_preserves_content_when_not_trailing_newline(self, tool, tmp_path):
|
||||
"""Deleting a word mid-line should not consume extra characters."""
|
||||
f = tmp_path / "a.py"
|
||||
f.write_text("hello world here\n", encoding="utf-8")
|
||||
result = await tool.execute(path=str(f), old_text="world ", new_text="")
|
||||
assert "Successfully" in result
|
||||
assert f.read_text() == "hello here\n"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Smart quote normalization
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestSmartQuoteNormalization:
|
||||
"""_find_match should handle curly ↔ straight quote fallback."""
|
||||
|
||||
def test_curly_double_quotes_match_straight(self):
|
||||
content = 'She said \u201chello\u201d to him'
|
||||
old_text = 'She said "hello" to him'
|
||||
match, count = _find_match(content, old_text)
|
||||
assert match is not None
|
||||
assert count == 1
|
||||
# Returned match should be the ORIGINAL content with curly quotes
|
||||
assert "\u201c" in match
|
||||
|
||||
def test_curly_single_quotes_match_straight(self):
|
||||
content = "it\u2019s a test"
|
||||
old_text = "it's a test"
|
||||
match, count = _find_match(content, old_text)
|
||||
assert match is not None
|
||||
assert count == 1
|
||||
assert "\u2019" in match
|
||||
|
||||
def test_straight_matches_curly_in_old_text(self):
|
||||
content = 'x = "hello"'
|
||||
old_text = 'x = \u201chello\u201d'
|
||||
match, count = _find_match(content, old_text)
|
||||
assert match is not None
|
||||
assert count == 1
|
||||
|
||||
def test_exact_match_still_preferred_over_quote_normalization(self):
|
||||
content = 'x = "hello"'
|
||||
old_text = 'x = "hello"'
|
||||
match, count = _find_match(content, old_text)
|
||||
assert match == old_text
|
||||
assert count == 1
|
||||
|
||||
|
||||
class TestQuoteStylePreservation:
|
||||
"""When quote-normalized matching occurs, replacement should preserve actual quote style."""
|
||||
|
||||
@pytest.fixture()
|
||||
def tool(self, tmp_path):
|
||||
return EditFileTool(workspace=tmp_path)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_replacement_preserves_curly_double_quotes(self, tool, tmp_path):
|
||||
f = tmp_path / "quotes.txt"
|
||||
f.write_text('message = “hello”\n', encoding="utf-8")
|
||||
result = await tool.execute(
|
||||
path=str(f),
|
||||
old_text='message = "hello"',
|
||||
new_text='message = "goodbye"',
|
||||
)
|
||||
assert "Successfully" in result
|
||||
assert f.read_text(encoding="utf-8") == 'message = “goodbye”\n'
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_replacement_preserves_curly_apostrophe(self, tool, tmp_path):
|
||||
f = tmp_path / "apostrophe.txt"
|
||||
f.write_text("it’s fine\n", encoding="utf-8")
|
||||
result = await tool.execute(
|
||||
path=str(f),
|
||||
old_text="it's fine",
|
||||
new_text="it's better",
|
||||
)
|
||||
assert "Successfully" in result
|
||||
assert f.read_text(encoding="utf-8") == "it’s better\n"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Indentation preservation
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestIndentationPreservation:
|
||||
"""Replacement should keep outer indentation when trim fallback matched."""
|
||||
|
||||
@pytest.fixture()
|
||||
def tool(self, tmp_path):
|
||||
return EditFileTool(workspace=tmp_path)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_trim_fallback_preserves_outer_indentation(self, tool, tmp_path):
|
||||
f = tmp_path / "indent.py"
|
||||
f.write_text(
|
||||
"if True:\n"
|
||||
" def foo():\n"
|
||||
" pass\n",
|
||||
encoding="utf-8",
|
||||
)
|
||||
result = await tool.execute(
|
||||
path=str(f),
|
||||
old_text="def foo():\n pass",
|
||||
new_text="def bar():\n return 1",
|
||||
)
|
||||
assert "Successfully" in result
|
||||
assert f.read_text(encoding="utf-8") == (
|
||||
"if True:\n"
|
||||
" def bar():\n"
|
||||
" return 1\n"
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Failure diagnostics
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestEditDiagnostics:
|
||||
"""Failure paths should offer actionable hints."""
|
||||
|
||||
@pytest.fixture()
|
||||
def tool(self, tmp_path):
|
||||
return EditFileTool(workspace=tmp_path)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_ambiguous_match_reports_candidate_lines(self, tool, tmp_path):
|
||||
f = tmp_path / "dup.py"
|
||||
f.write_text("aaa\nbbb\naaa\nbbb\n", encoding="utf-8")
|
||||
result = await tool.execute(path=str(f), old_text="aaa\nbbb", new_text="xxx")
|
||||
assert "appears 2 times" in result.lower()
|
||||
assert "line 1" in result.lower()
|
||||
assert "line 3" in result.lower()
|
||||
assert "replace_all=true" in result
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_not_found_reports_whitespace_hint(self, tool, tmp_path):
|
||||
f = tmp_path / "space.py"
|
||||
f.write_text("value = 1\n", encoding="utf-8")
|
||||
result = await tool.execute(path=str(f), old_text="value = 1", new_text="value = 2")
|
||||
assert "Error" in result
|
||||
assert "whitespace" in result.lower()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_not_found_reports_case_hint(self, tool, tmp_path):
|
||||
f = tmp_path / "case.py"
|
||||
f.write_text("HelloWorld\n", encoding="utf-8")
|
||||
result = await tool.execute(path=str(f), old_text="helloworld", new_text="goodbye")
|
||||
assert "Error" in result
|
||||
assert "letter case differs" in result.lower()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Advanced fallback replacement behavior
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestAdvancedReplaceAll:
|
||||
"""replace_all should work correctly for fallback-based matches too."""
|
||||
|
||||
@pytest.fixture()
|
||||
def tool(self, tmp_path):
|
||||
return EditFileTool(workspace=tmp_path)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_replace_all_preserves_each_match_indentation(self, tool, tmp_path):
|
||||
f = tmp_path / "indent_multi.py"
|
||||
f.write_text(
|
||||
"if a:\n"
|
||||
" def foo():\n"
|
||||
" pass\n"
|
||||
"if b:\n"
|
||||
" def foo():\n"
|
||||
" pass\n",
|
||||
encoding="utf-8",
|
||||
)
|
||||
result = await tool.execute(
|
||||
path=str(f),
|
||||
old_text="def foo():\n pass",
|
||||
new_text="def bar():\n return 1",
|
||||
replace_all=True,
|
||||
)
|
||||
assert "Successfully" in result
|
||||
assert f.read_text(encoding="utf-8") == (
|
||||
"if a:\n"
|
||||
" def bar():\n"
|
||||
" return 1\n"
|
||||
"if b:\n"
|
||||
" def bar():\n"
|
||||
" return 1\n"
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_trim_and_quote_fallback_match_succeeds(self, tool, tmp_path):
|
||||
f = tmp_path / "quote_indent.py"
|
||||
f.write_text(" message = “hello”\n", encoding="utf-8")
|
||||
result = await tool.execute(
|
||||
path=str(f),
|
||||
old_text='message = "hello"',
|
||||
new_text='message = "goodbye"',
|
||||
)
|
||||
assert "Successfully" in result
|
||||
assert f.read_text(encoding="utf-8") == " message = “goodbye”\n"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Advanced fallback replacement behavior
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestAdvancedReplaceAll:
|
||||
"""replace_all should work correctly for fallback-based matches too."""
|
||||
|
||||
@pytest.fixture()
|
||||
def tool(self, tmp_path):
|
||||
return EditFileTool(workspace=tmp_path)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_replace_all_preserves_each_match_indentation(self, tool, tmp_path):
|
||||
f = tmp_path / "indent_multi.py"
|
||||
f.write_text(
|
||||
"if a:\n"
|
||||
" def foo():\n"
|
||||
" pass\n"
|
||||
"if b:\n"
|
||||
" def foo():\n"
|
||||
" pass\n",
|
||||
encoding="utf-8",
|
||||
)
|
||||
result = await tool.execute(
|
||||
path=str(f),
|
||||
old_text="def foo():\n pass",
|
||||
new_text="def bar():\n return 1",
|
||||
replace_all=True,
|
||||
)
|
||||
assert "Successfully" in result
|
||||
assert f.read_text(encoding="utf-8") == (
|
||||
"if a:\n"
|
||||
" def bar():\n"
|
||||
" return 1\n"
|
||||
"if b:\n"
|
||||
" def bar():\n"
|
||||
" return 1\n"
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_trim_and_quote_fallback_match_succeeds(self, tool, tmp_path):
|
||||
f = tmp_path / "quote_indent.py"
|
||||
f.write_text(" message = “hello”\n", encoding="utf-8")
|
||||
result = await tool.execute(
|
||||
path=str(f),
|
||||
old_text='message = "hello"',
|
||||
new_text='message = "goodbye"',
|
||||
)
|
||||
assert "Successfully" in result
|
||||
assert f.read_text(encoding="utf-8") == " message = “goodbye”\n"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Trailing whitespace stripping on new_text
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestTrailingWhitespaceStrip:
|
||||
"""new_text trailing whitespace should be stripped (except .md files)."""
|
||||
|
||||
@pytest.fixture()
|
||||
def tool(self, tmp_path):
|
||||
return EditFileTool(workspace=tmp_path)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_strips_trailing_whitespace_from_new_text(self, tool, tmp_path):
|
||||
f = tmp_path / "a.py"
|
||||
f.write_text("x = 1\n", encoding="utf-8")
|
||||
result = await tool.execute(
|
||||
path=str(f), old_text="x = 1", new_text="x = 2 \ny = 3 ",
|
||||
)
|
||||
assert "Successfully" in result
|
||||
content = f.read_text()
|
||||
assert "x = 2\ny = 3\n" == content
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_preserves_trailing_whitespace_in_markdown(self, tool, tmp_path):
|
||||
f = tmp_path / "doc.md"
|
||||
f.write_text("# Title\n", encoding="utf-8")
|
||||
# Markdown uses trailing double-space for line breaks
|
||||
result = await tool.execute(
|
||||
path=str(f), old_text="# Title", new_text="# Title \nSubtitle ",
|
||||
)
|
||||
assert "Successfully" in result
|
||||
content = f.read_text()
|
||||
# Trailing spaces should be preserved for markdown
|
||||
assert "Title " in content
|
||||
assert "Subtitle " in content
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# File size protection
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestFileSizeProtection:
|
||||
"""Editing extremely large files should be rejected."""
|
||||
|
||||
@pytest.fixture()
|
||||
def tool(self, tmp_path):
|
||||
return EditFileTool(workspace=tmp_path)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_rejects_file_over_size_limit(self, tool, tmp_path):
|
||||
f = tmp_path / "huge.txt"
|
||||
f.write_text("x", encoding="utf-8")
|
||||
# Monkey-patch the file size check by creating a stat mock
|
||||
original_stat = f.stat
|
||||
|
||||
class FakeStat:
|
||||
def __init__(self, real_stat):
|
||||
self._real = real_stat
|
||||
|
||||
def __getattr__(self, name):
|
||||
return getattr(self._real, name)
|
||||
|
||||
@property
|
||||
def st_size(self):
|
||||
return 2 * 1024 * 1024 * 1024 # 2 GiB
|
||||
|
||||
import unittest.mock
|
||||
with unittest.mock.patch.object(type(f), 'stat', return_value=FakeStat(f.stat())):
|
||||
result = await tool.execute(path=str(f), old_text="x", new_text="y")
|
||||
assert "Error" in result
|
||||
assert "too large" in result.lower() or "size" in result.lower()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Stale detection with content-equality fallback
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestStaleDetectionContentFallback:
|
||||
"""When mtime changed but file content is unchanged, edit should proceed without warning."""
|
||||
|
||||
@pytest.fixture()
|
||||
def read_tool(self, tmp_path):
|
||||
return ReadFileTool(workspace=tmp_path)
|
||||
|
||||
@pytest.fixture()
|
||||
def edit_tool(self, tmp_path):
|
||||
return EditFileTool(workspace=tmp_path)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_mtime_bump_same_content_no_warning(self, read_tool, edit_tool, tmp_path):
|
||||
f = tmp_path / "a.py"
|
||||
f.write_text("hello world", encoding="utf-8")
|
||||
await read_tool.execute(path=str(f))
|
||||
|
||||
# Touch the file to bump mtime without changing content
|
||||
time.sleep(0.05)
|
||||
original_content = f.read_text()
|
||||
f.write_text(original_content, encoding="utf-8")
|
||||
|
||||
result = await edit_tool.execute(path=str(f), old_text="world", new_text="earth")
|
||||
assert "Successfully" in result
|
||||
# Should NOT warn about modification since content is the same
|
||||
assert "modified" not in result.lower()
|
||||
@@ -0,0 +1,152 @@
|
||||
"""Tests for EditFileTool enhancements: read-before-edit tracking, path suggestions,
|
||||
.ipynb detection, and create-file semantics."""
|
||||
|
||||
import pytest
|
||||
|
||||
from nanobot.agent.tools.filesystem import EditFileTool, ReadFileTool, WriteFileTool
|
||||
from nanobot.agent.tools import file_state
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _clear_file_state():
|
||||
"""Reset global read-state between tests."""
|
||||
file_state.clear()
|
||||
yield
|
||||
file_state.clear()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Read-before-edit tracking
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class TestEditReadTracking:
|
||||
"""edit_file should warn when file hasn't been read first."""
|
||||
|
||||
@pytest.fixture()
|
||||
def read_tool(self, tmp_path):
|
||||
return ReadFileTool(workspace=tmp_path)
|
||||
|
||||
@pytest.fixture()
|
||||
def edit_tool(self, tmp_path):
|
||||
return EditFileTool(workspace=tmp_path)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_edit_warns_if_file_not_read_first(self, edit_tool, tmp_path):
|
||||
f = tmp_path / "a.py"
|
||||
f.write_text("hello world", encoding="utf-8")
|
||||
result = await edit_tool.execute(path=str(f), old_text="world", new_text="earth")
|
||||
# Should still succeed but include a warning
|
||||
assert "Successfully" in result
|
||||
assert "not been read" in result.lower() or "warning" in result.lower()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_edit_succeeds_cleanly_after_read(self, read_tool, edit_tool, tmp_path):
|
||||
f = tmp_path / "a.py"
|
||||
f.write_text("hello world", encoding="utf-8")
|
||||
await read_tool.execute(path=str(f))
|
||||
result = await edit_tool.execute(path=str(f), old_text="world", new_text="earth")
|
||||
assert "Successfully" in result
|
||||
# No warning when file was read first
|
||||
assert "not been read" not in result.lower()
|
||||
assert f.read_text() == "hello earth"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_edit_warns_if_file_modified_since_read(self, read_tool, edit_tool, tmp_path):
|
||||
f = tmp_path / "a.py"
|
||||
f.write_text("hello world", encoding="utf-8")
|
||||
await read_tool.execute(path=str(f))
|
||||
# External modification
|
||||
f.write_text("hello universe", encoding="utf-8")
|
||||
result = await edit_tool.execute(path=str(f), old_text="universe", new_text="earth")
|
||||
assert "Successfully" in result
|
||||
assert "modified" in result.lower() or "warning" in result.lower()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Create-file semantics
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class TestEditCreateFile:
|
||||
"""edit_file with old_text='' creates new file if not exists."""
|
||||
|
||||
@pytest.fixture()
|
||||
def tool(self, tmp_path):
|
||||
return EditFileTool(workspace=tmp_path)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_create_new_file_with_empty_old_text(self, tool, tmp_path):
|
||||
f = tmp_path / "subdir" / "new.py"
|
||||
result = await tool.execute(path=str(f), old_text="", new_text="print('hi')")
|
||||
assert "created" in result.lower() or "Successfully" in result
|
||||
assert f.exists()
|
||||
assert f.read_text() == "print('hi')"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_create_fails_if_file_already_exists_and_not_empty(self, tool, tmp_path):
|
||||
f = tmp_path / "existing.py"
|
||||
f.write_text("existing content", encoding="utf-8")
|
||||
result = await tool.execute(path=str(f), old_text="", new_text="new content")
|
||||
assert "Error" in result or "already exists" in result.lower()
|
||||
# File should be unchanged
|
||||
assert f.read_text() == "existing content"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_create_succeeds_if_file_exists_but_empty(self, tool, tmp_path):
|
||||
f = tmp_path / "empty.py"
|
||||
f.write_text("", encoding="utf-8")
|
||||
result = await tool.execute(path=str(f), old_text="", new_text="print('hi')")
|
||||
assert "Successfully" in result
|
||||
assert f.read_text() == "print('hi')"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# .ipynb detection
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class TestEditIpynbDetection:
|
||||
"""edit_file should refuse .ipynb and suggest notebook_edit."""
|
||||
|
||||
@pytest.fixture()
|
||||
def tool(self, tmp_path):
|
||||
return EditFileTool(workspace=tmp_path)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_ipynb_rejected_with_suggestion(self, tool, tmp_path):
|
||||
f = tmp_path / "analysis.ipynb"
|
||||
f.write_text('{"cells": []}', encoding="utf-8")
|
||||
result = await tool.execute(path=str(f), old_text="x", new_text="y")
|
||||
assert "notebook" in result.lower()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Path suggestion on not-found
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class TestEditPathSuggestion:
|
||||
"""edit_file should suggest similar paths on not-found."""
|
||||
|
||||
@pytest.fixture()
|
||||
def tool(self, tmp_path):
|
||||
return EditFileTool(workspace=tmp_path)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_suggests_similar_filename(self, tool, tmp_path):
|
||||
f = tmp_path / "config.py"
|
||||
f.write_text("x = 1", encoding="utf-8")
|
||||
# Typo: conifg.py
|
||||
result = await tool.execute(
|
||||
path=str(tmp_path / "conifg.py"), old_text="x = 1", new_text="x = 2",
|
||||
)
|
||||
assert "Error" in result
|
||||
assert "config.py" in result
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_shows_cwd_in_error(self, tool, tmp_path):
|
||||
result = await tool.execute(
|
||||
path=str(tmp_path / "nonexistent.py"), old_text="a", new_text="b",
|
||||
)
|
||||
assert "Error" in result
|
||||
@@ -0,0 +1,147 @@
|
||||
"""Tests for NotebookEditTool — Jupyter .ipynb editing."""
|
||||
|
||||
import json
|
||||
|
||||
import pytest
|
||||
|
||||
from nanobot.agent.tools.notebook import NotebookEditTool
|
||||
|
||||
|
||||
def _make_notebook(cells: list[dict] | None = None, nbformat: int = 4, nbformat_minor: int = 5) -> dict:
|
||||
"""Build a minimal valid .ipynb structure."""
|
||||
return {
|
||||
"nbformat": nbformat,
|
||||
"nbformat_minor": nbformat_minor,
|
||||
"metadata": {"kernelspec": {"display_name": "Python 3", "language": "python", "name": "python3"}},
|
||||
"cells": cells or [],
|
||||
}
|
||||
|
||||
|
||||
def _code_cell(source: str, cell_id: str | None = None) -> dict:
|
||||
cell = {"cell_type": "code", "source": source, "metadata": {}, "outputs": [], "execution_count": None}
|
||||
if cell_id:
|
||||
cell["id"] = cell_id
|
||||
return cell
|
||||
|
||||
|
||||
def _md_cell(source: str, cell_id: str | None = None) -> dict:
|
||||
cell = {"cell_type": "markdown", "source": source, "metadata": {}}
|
||||
if cell_id:
|
||||
cell["id"] = cell_id
|
||||
return cell
|
||||
|
||||
|
||||
def _write_nb(tmp_path, name: str, nb: dict) -> str:
|
||||
p = tmp_path / name
|
||||
p.write_text(json.dumps(nb), encoding="utf-8")
|
||||
return str(p)
|
||||
|
||||
|
||||
class TestNotebookEdit:
|
||||
|
||||
@pytest.fixture()
|
||||
def tool(self, tmp_path):
|
||||
return NotebookEditTool(workspace=tmp_path)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_replace_cell_content(self, tool, tmp_path):
|
||||
nb = _make_notebook([_code_cell("print('hello')"), _code_cell("x = 1")])
|
||||
path = _write_nb(tmp_path, "test.ipynb", nb)
|
||||
result = await tool.execute(path=path, cell_index=0, new_source="print('world')")
|
||||
assert "Successfully" in result
|
||||
saved = json.loads((tmp_path / "test.ipynb").read_text())
|
||||
assert saved["cells"][0]["source"] == "print('world')"
|
||||
assert saved["cells"][1]["source"] == "x = 1"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_insert_cell_after_target(self, tool, tmp_path):
|
||||
nb = _make_notebook([_code_cell("cell 0"), _code_cell("cell 1")])
|
||||
path = _write_nb(tmp_path, "test.ipynb", nb)
|
||||
result = await tool.execute(path=path, cell_index=0, new_source="inserted", edit_mode="insert")
|
||||
assert "Successfully" in result
|
||||
saved = json.loads((tmp_path / "test.ipynb").read_text())
|
||||
assert len(saved["cells"]) == 3
|
||||
assert saved["cells"][0]["source"] == "cell 0"
|
||||
assert saved["cells"][1]["source"] == "inserted"
|
||||
assert saved["cells"][2]["source"] == "cell 1"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_delete_cell(self, tool, tmp_path):
|
||||
nb = _make_notebook([_code_cell("A"), _code_cell("B"), _code_cell("C")])
|
||||
path = _write_nb(tmp_path, "test.ipynb", nb)
|
||||
result = await tool.execute(path=path, cell_index=1, edit_mode="delete")
|
||||
assert "Successfully" in result
|
||||
saved = json.loads((tmp_path / "test.ipynb").read_text())
|
||||
assert len(saved["cells"]) == 2
|
||||
assert saved["cells"][0]["source"] == "A"
|
||||
assert saved["cells"][1]["source"] == "C"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_create_new_notebook_from_scratch(self, tool, tmp_path):
|
||||
path = str(tmp_path / "new.ipynb")
|
||||
result = await tool.execute(path=path, cell_index=0, new_source="# Hello", edit_mode="insert", cell_type="markdown")
|
||||
assert "Successfully" in result or "created" in result.lower()
|
||||
saved = json.loads((tmp_path / "new.ipynb").read_text())
|
||||
assert saved["nbformat"] == 4
|
||||
assert len(saved["cells"]) == 1
|
||||
assert saved["cells"][0]["cell_type"] == "markdown"
|
||||
assert saved["cells"][0]["source"] == "# Hello"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_invalid_cell_index_error(self, tool, tmp_path):
|
||||
nb = _make_notebook([_code_cell("only cell")])
|
||||
path = _write_nb(tmp_path, "test.ipynb", nb)
|
||||
result = await tool.execute(path=path, cell_index=5, new_source="x")
|
||||
assert "Error" in result
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_non_ipynb_rejected(self, tool, tmp_path):
|
||||
f = tmp_path / "script.py"
|
||||
f.write_text("pass")
|
||||
result = await tool.execute(path=str(f), cell_index=0, new_source="x")
|
||||
assert "Error" in result
|
||||
assert ".ipynb" in result
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_preserves_metadata_and_outputs(self, tool, tmp_path):
|
||||
cell = _code_cell("old")
|
||||
cell["outputs"] = [{"output_type": "stream", "text": "hello\n"}]
|
||||
cell["execution_count"] = 42
|
||||
nb = _make_notebook([cell])
|
||||
path = _write_nb(tmp_path, "test.ipynb", nb)
|
||||
await tool.execute(path=path, cell_index=0, new_source="new")
|
||||
saved = json.loads((tmp_path / "test.ipynb").read_text())
|
||||
assert saved["metadata"]["kernelspec"]["language"] == "python"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_nbformat_45_generates_cell_id(self, tool, tmp_path):
|
||||
nb = _make_notebook([], nbformat_minor=5)
|
||||
path = _write_nb(tmp_path, "test.ipynb", nb)
|
||||
await tool.execute(path=path, cell_index=0, new_source="x = 1", edit_mode="insert")
|
||||
saved = json.loads((tmp_path / "test.ipynb").read_text())
|
||||
assert "id" in saved["cells"][0]
|
||||
assert len(saved["cells"][0]["id"]) > 0
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_insert_with_cell_type_markdown(self, tool, tmp_path):
|
||||
nb = _make_notebook([_code_cell("code")])
|
||||
path = _write_nb(tmp_path, "test.ipynb", nb)
|
||||
await tool.execute(path=path, cell_index=0, new_source="# Title", edit_mode="insert", cell_type="markdown")
|
||||
saved = json.loads((tmp_path / "test.ipynb").read_text())
|
||||
assert saved["cells"][1]["cell_type"] == "markdown"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_invalid_edit_mode_rejected(self, tool, tmp_path):
|
||||
nb = _make_notebook([_code_cell("code")])
|
||||
path = _write_nb(tmp_path, "test.ipynb", nb)
|
||||
result = await tool.execute(path=path, cell_index=0, new_source="x", edit_mode="replcae")
|
||||
assert "Error" in result
|
||||
assert "edit_mode" in result
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_invalid_cell_type_rejected(self, tool, tmp_path):
|
||||
nb = _make_notebook([_code_cell("code")])
|
||||
path = _write_nb(tmp_path, "test.ipynb", nb)
|
||||
result = await tool.execute(path=path, cell_index=0, new_source="x", cell_type="raw")
|
||||
assert "Error" in result
|
||||
assert "cell_type" in result
|
||||
@@ -0,0 +1,180 @@
|
||||
"""Tests for ReadFileTool enhancements: description fix, read dedup, PDF support, device blacklist."""
|
||||
|
||||
import pytest
|
||||
|
||||
from nanobot.agent.tools.filesystem import ReadFileTool, WriteFileTool
|
||||
from nanobot.agent.tools import file_state
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _clear_file_state():
|
||||
file_state.clear()
|
||||
yield
|
||||
file_state.clear()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Description fix
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class TestReadDescriptionFix:
|
||||
|
||||
def test_description_mentions_image_support(self):
|
||||
tool = ReadFileTool()
|
||||
assert "image" in tool.description.lower()
|
||||
|
||||
def test_description_no_longer_says_cannot_read_images(self):
|
||||
tool = ReadFileTool()
|
||||
assert "cannot read binary files or images" not in tool.description.lower()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Read deduplication
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class TestReadDedup:
|
||||
"""Same file + same offset/limit + unchanged mtime -> short stub."""
|
||||
|
||||
@pytest.fixture()
|
||||
def tool(self, tmp_path):
|
||||
return ReadFileTool(workspace=tmp_path)
|
||||
|
||||
@pytest.fixture()
|
||||
def write_tool(self, tmp_path):
|
||||
return WriteFileTool(workspace=tmp_path)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_second_read_returns_unchanged_stub(self, tool, tmp_path):
|
||||
f = tmp_path / "data.txt"
|
||||
f.write_text("\n".join(f"line {i}" for i in range(100)), encoding="utf-8")
|
||||
first = await tool.execute(path=str(f))
|
||||
assert "line 0" in first
|
||||
second = await tool.execute(path=str(f))
|
||||
assert "unchanged" in second.lower()
|
||||
# Stub should not contain file content
|
||||
assert "line 0" not in second
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_read_after_external_modification_returns_full(self, tool, tmp_path):
|
||||
f = tmp_path / "data.txt"
|
||||
f.write_text("original", encoding="utf-8")
|
||||
await tool.execute(path=str(f))
|
||||
# Modify the file externally
|
||||
f.write_text("modified content", encoding="utf-8")
|
||||
second = await tool.execute(path=str(f))
|
||||
assert "modified content" in second
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_different_offset_returns_full(self, tool, tmp_path):
|
||||
f = tmp_path / "data.txt"
|
||||
f.write_text("\n".join(f"line {i}" for i in range(1, 21)), encoding="utf-8")
|
||||
await tool.execute(path=str(f), offset=1, limit=5)
|
||||
second = await tool.execute(path=str(f), offset=6, limit=5)
|
||||
# Different offset → full read, not stub
|
||||
assert "line 6" in second
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_first_read_after_write_returns_full_content(self, tool, write_tool, tmp_path):
|
||||
f = tmp_path / "fresh.txt"
|
||||
result = await write_tool.execute(path=str(f), content="hello")
|
||||
assert "Successfully" in result
|
||||
read_result = await tool.execute(path=str(f))
|
||||
assert "hello" in read_result
|
||||
assert "unchanged" not in read_result.lower()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_dedup_does_not_apply_to_images(self, tool, tmp_path):
|
||||
f = tmp_path / "img.png"
|
||||
f.write_bytes(b"\x89PNG\r\n\x1a\nfake-png-data")
|
||||
first = await tool.execute(path=str(f))
|
||||
assert isinstance(first, list)
|
||||
second = await tool.execute(path=str(f))
|
||||
# Images should always return full content blocks, not a stub
|
||||
assert isinstance(second, list)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# PDF support
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class TestReadPdf:
|
||||
|
||||
@pytest.fixture()
|
||||
def tool(self, tmp_path):
|
||||
return ReadFileTool(workspace=tmp_path)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pdf_returns_text_content(self, tool, tmp_path):
|
||||
fitz = pytest.importorskip("fitz")
|
||||
pdf_path = tmp_path / "test.pdf"
|
||||
doc = fitz.open()
|
||||
page = doc.new_page()
|
||||
page.insert_text((72, 72), "Hello PDF World")
|
||||
doc.save(str(pdf_path))
|
||||
doc.close()
|
||||
|
||||
result = await tool.execute(path=str(pdf_path))
|
||||
assert "Hello PDF World" in result
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pdf_pages_parameter(self, tool, tmp_path):
|
||||
fitz = pytest.importorskip("fitz")
|
||||
pdf_path = tmp_path / "multi.pdf"
|
||||
doc = fitz.open()
|
||||
for i in range(5):
|
||||
page = doc.new_page()
|
||||
page.insert_text((72, 72), f"Page {i + 1} content")
|
||||
doc.save(str(pdf_path))
|
||||
doc.close()
|
||||
|
||||
result = await tool.execute(path=str(pdf_path), pages="2-3")
|
||||
assert "Page 2 content" in result
|
||||
assert "Page 3 content" in result
|
||||
assert "Page 1 content" not in result
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pdf_file_not_found_error(self, tool, tmp_path):
|
||||
result = await tool.execute(path=str(tmp_path / "nope.pdf"))
|
||||
assert "Error" in result
|
||||
assert "not found" in result
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Device path blacklist
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class TestReadDeviceBlacklist:
|
||||
|
||||
@pytest.fixture()
|
||||
def tool(self):
|
||||
return ReadFileTool()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_dev_random_blocked(self, tool):
|
||||
result = await tool.execute(path="/dev/random")
|
||||
assert "Error" in result
|
||||
assert "blocked" in result.lower() or "device" in result.lower()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_dev_urandom_blocked(self, tool):
|
||||
result = await tool.execute(path="/dev/urandom")
|
||||
assert "Error" in result
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_dev_zero_blocked(self, tool):
|
||||
result = await tool.execute(path="/dev/zero")
|
||||
assert "Error" in result
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_proc_fd_blocked(self, tool):
|
||||
result = await tool.execute(path="/proc/self/fd/0")
|
||||
assert "Error" in result
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_symlink_to_dev_zero_blocked(self, tmp_path):
|
||||
tool = ReadFileTool(workspace=tmp_path)
|
||||
link = tmp_path / "zero-link"
|
||||
link.symlink_to("/dev/zero")
|
||||
result = await tool.execute(path=str(link))
|
||||
assert "Error" in result
|
||||
assert "blocked" in result.lower() or "device" in result.lower()
|
||||
Reference in New Issue
Block a user