140 lines
4.5 KiB
Python
140 lines
4.5 KiB
Python
from __future__ import annotations
|
|
|
|
from unittest.mock import AsyncMock, patch
|
|
|
|
import pytest
|
|
|
|
from app.agents.folder_agent import (
|
|
read_project_folder_file,
|
|
search_project_folder_file,
|
|
)
|
|
|
|
pytestmark = pytest.mark.asyncio
|
|
|
|
|
|
async def test_happy_path():
|
|
with patch(
|
|
"app.agents.folder_agent.execute_on_client",
|
|
new=AsyncMock(return_value={"content": "file body", "kind": "text", "totalSize": 9}),
|
|
):
|
|
out = await read_project_folder_file.ainvoke({"project_id": "p1", "relative_path": "docs/x.md"})
|
|
assert "file body" in out
|
|
assert "kind=text" in out
|
|
|
|
|
|
async def test_traversal_rejected():
|
|
out = await read_project_folder_file.ainvoke({"project_id": "p1", "relative_path": "../../etc/passwd"})
|
|
assert out == "Access denied"
|
|
|
|
|
|
async def test_absolute_path_rejected():
|
|
out = await read_project_folder_file.ainvoke({"project_id": "p1", "relative_path": "C:\\Windows\\foo"})
|
|
assert out == "Access denied"
|
|
|
|
|
|
async def test_missing_file():
|
|
with patch(
|
|
"app.agents.folder_agent.execute_on_client",
|
|
new=AsyncMock(return_value={"content": "", "kind": "missing", "totalSize": 0}),
|
|
):
|
|
out = await read_project_folder_file.ainvoke({"project_id": "p1", "relative_path": "ghost.md"})
|
|
assert "not found" in out.lower()
|
|
|
|
|
|
async def test_pagination_signals_more_available():
|
|
# Electron returned the first slice, totalSize larger than slice length.
|
|
with patch(
|
|
"app.agents.folder_agent.execute_on_client",
|
|
new=AsyncMock(return_value={"content": "first chunk", "kind": "text", "totalSize": 1000}),
|
|
):
|
|
out = await read_project_folder_file.ainvoke({
|
|
"project_id": "p1",
|
|
"relative_path": "big.txt",
|
|
"offset": 0,
|
|
"length": 11,
|
|
})
|
|
assert "first chunk" in out
|
|
assert "More content available" in out
|
|
assert "offset=11" in out
|
|
|
|
|
|
async def test_pdf_extracted_then_sliced(monkeypatch):
|
|
from app.agents import folder_agent
|
|
monkeypatch.setattr(folder_agent, "_extract_pdf_text", lambda b: "ABC " * 100)
|
|
with patch(
|
|
"app.agents.folder_agent.execute_on_client",
|
|
new=AsyncMock(return_value={"content": "JVBERi0xLg==", "kind": "pdf", "totalSize": 12}),
|
|
):
|
|
out = await read_project_folder_file.ainvoke({
|
|
"project_id": "p1",
|
|
"relative_path": "doc.pdf",
|
|
"offset": 0,
|
|
"length": 8,
|
|
})
|
|
assert "kind=pdf" in out
|
|
assert "ABC ABC " in out
|
|
assert "More content available" in out
|
|
|
|
|
|
async def test_image_returns_placeholder():
|
|
with patch(
|
|
"app.agents.folder_agent.execute_on_client",
|
|
new=AsyncMock(return_value={"content": "iVBORw0K", "kind": "image", "totalSize": 1024}),
|
|
):
|
|
out = await read_project_folder_file.ainvoke({"project_id": "p1", "relative_path": "logo.png"})
|
|
assert "image" in out.lower()
|
|
|
|
|
|
async def test_search_finds_match_with_context():
|
|
body = "alpha\nbeta\nthe needle is here\ngamma\ndelta"
|
|
with patch(
|
|
"app.agents.folder_agent.execute_on_client",
|
|
new=AsyncMock(return_value={"content": body, "kind": "text", "totalSize": len(body)}),
|
|
):
|
|
out = await search_project_folder_file.ainvoke({
|
|
"project_id": "p1",
|
|
"relative_path": "log.txt",
|
|
"query": "needle",
|
|
"context_lines": 1,
|
|
})
|
|
assert "needle" in out
|
|
assert "matches=1" in out
|
|
# Context lines included
|
|
assert "beta" in out
|
|
assert "gamma" in out
|
|
|
|
|
|
async def test_search_no_match():
|
|
with patch(
|
|
"app.agents.folder_agent.execute_on_client",
|
|
new=AsyncMock(return_value={"content": "nothing here", "kind": "text", "totalSize": 12}),
|
|
):
|
|
out = await search_project_folder_file.ainvoke({
|
|
"project_id": "p1",
|
|
"relative_path": "x.txt",
|
|
"query": "zzz",
|
|
})
|
|
assert "No matches" in out
|
|
|
|
|
|
async def test_search_rejects_traversal():
|
|
out = await search_project_folder_file.ainvoke({
|
|
"project_id": "p1",
|
|
"relative_path": "../etc/passwd",
|
|
"query": "root",
|
|
})
|
|
assert out == "Access denied"
|
|
|
|
|
|
async def test_search_image_rejected():
|
|
with patch(
|
|
"app.agents.folder_agent.execute_on_client",
|
|
new=AsyncMock(return_value={"content": "b64data", "kind": "image", "totalSize": 100}),
|
|
):
|
|
out = await search_project_folder_file.ainvoke({
|
|
"project_id": "p1",
|
|
"relative_path": "logo.png",
|
|
"query": "anything",
|
|
})
|
|
assert "Cannot search" in out
|