97 lines
3.2 KiB
Python
97 lines
3.2 KiB
Python
from __future__ import annotations
|
|
|
|
import tempfile
|
|
import unittest
|
|
from pathlib import Path
|
|
from types import SimpleNamespace
|
|
from unittest.mock import AsyncMock
|
|
|
|
from micro_agents_tools import FilesystemTools, PythonTools
|
|
|
|
from app.agent_tools import AgentToolManager, TelegramTools
|
|
|
|
|
|
class AgentToolPolicyTests(unittest.IsolatedAsyncioTestCase):
|
|
async def test_filesystem_and_python_tools_are_owner_only(self) -> None:
|
|
telegram_tools = TelegramTools(
|
|
SimpleNamespace(),
|
|
99,
|
|
base_dir=Path.cwd(),
|
|
).tools()
|
|
manager = AgentToolManager(
|
|
[*FilesystemTools().tools(), *PythonTools().tools(), *telegram_tools]
|
|
)
|
|
|
|
guest_names = {
|
|
tool.name for tool in await manager.available_tools({"is_owner": False})
|
|
}
|
|
owner_names = {
|
|
tool.name for tool in await manager.available_tools({"is_owner": True})
|
|
}
|
|
|
|
self.assertEqual(guest_names, {"send_message", "send_image", "send_file"})
|
|
self.assertIn("read_file", owner_names)
|
|
self.assertIn("write_file", owner_names)
|
|
self.assertIn("execute_python", owner_names)
|
|
|
|
|
|
class TelegramToolsTests(unittest.IsolatedAsyncioTestCase):
|
|
async def test_sends_text_to_current_chat_as_reply_to_command(self) -> None:
|
|
client = SimpleNamespace(
|
|
send_message=AsyncMock(return_value=SimpleNamespace(id=123))
|
|
)
|
|
tools = TelegramTools(
|
|
client,
|
|
99,
|
|
base_dir=Path.cwd(),
|
|
reply_to_message_id=7,
|
|
)
|
|
|
|
result = await tools.send_message("hello")
|
|
|
|
client.send_message.assert_awaited_once_with(
|
|
99,
|
|
"hello",
|
|
reply_to_message_id=7,
|
|
)
|
|
self.assertIn("message_id=123", result)
|
|
|
|
async def test_sends_image_and_file_from_workspace(self) -> None:
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
base_dir = Path(temp_dir)
|
|
image = base_dir / "photo.jpg"
|
|
document = base_dir / "report.txt"
|
|
image.write_bytes(b"image")
|
|
document.write_text("report", encoding="utf-8")
|
|
client = SimpleNamespace(
|
|
send_photo=AsyncMock(return_value=SimpleNamespace(id=1)),
|
|
send_document=AsyncMock(return_value=SimpleNamespace(id=2)),
|
|
)
|
|
tools = TelegramTools(client, 99, base_dir=base_dir)
|
|
|
|
await tools.send_image("photo.jpg", "caption")
|
|
await tools.send_file("report.txt")
|
|
|
|
client.send_photo.assert_awaited_once_with(
|
|
99,
|
|
str(image),
|
|
caption="caption",
|
|
reply_to_message_id=None,
|
|
)
|
|
client.send_document.assert_awaited_once_with(
|
|
99,
|
|
str(document),
|
|
caption=None,
|
|
reply_to_message_id=None,
|
|
)
|
|
|
|
async def test_rejects_path_outside_workspace(self) -> None:
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
tools = TelegramTools(SimpleNamespace(), 99, base_dir=Path(temp_dir))
|
|
with self.assertRaisesRegex(ValueError, "Absolute paths"):
|
|
await tools.send_file(str(Path(temp_dir).resolve() / "secret.txt"))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|