feat: rebuild AI commands and migrate config to TOML
This commit is contained in:
Azimkin
306
tests/test_ai_command.py
Normal file
306
tests/test_ai_command.py
Normal file
@@ -0,0 +1,306 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import tempfile
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
from app.agentic import (
|
||||
AgentSessionHandle,
|
||||
AgentSessionStore,
|
||||
AgentSyntaxError,
|
||||
_build_context_messages,
|
||||
_download_message_media,
|
||||
render_prompt_template,
|
||||
parse_agent_invocation,
|
||||
run_agent_turn,
|
||||
strip_prompt_comments,
|
||||
)
|
||||
from app.commands.builtins import ai, command_registry
|
||||
from app.config import AiConfig, AiProviderConfig
|
||||
|
||||
|
||||
class AgentInvocationTests(unittest.TestCase):
|
||||
def test_prompt_comments_are_removed_without_breaking_urls(self) -> None:
|
||||
value = strip_prompt_comments(
|
||||
"// first\nkeep https://example.com // trailing\n/* block\ncomment */\nlast"
|
||||
)
|
||||
|
||||
self.assertNotIn("first", value)
|
||||
self.assertNotIn("trailing", value)
|
||||
self.assertNotIn("block", value)
|
||||
self.assertIn("https://example.com", value)
|
||||
self.assertIn("last", value)
|
||||
|
||||
def test_prompt_template_replaces_nested_values_and_text(self) -> None:
|
||||
rendered = render_prompt_template(
|
||||
"caller={{caller_user.id}} owner={owner.name} text={{text}}",
|
||||
{
|
||||
"caller_user": {"id": 42},
|
||||
"owner": {"name": "Owner"},
|
||||
"text": "old",
|
||||
"source_text": "new",
|
||||
},
|
||||
)
|
||||
|
||||
self.assertEqual(rendered, "caller=42 owner=Owner text=new")
|
||||
|
||||
def test_parses_context_provider_and_thinking_flags(self) -> None:
|
||||
invocation = parse_agent_invocation("12 .ds .th explain this", AiConfig())
|
||||
|
||||
self.assertEqual(invocation.context_messages, 12)
|
||||
self.assertTrue(invocation.use_deepseek)
|
||||
self.assertTrue(invocation.thinking)
|
||||
self.assertEqual(invocation.prompt, "explain this")
|
||||
|
||||
def test_ignores_unknown_leading_dot_arguments(self) -> None:
|
||||
invocation = parse_agent_invocation(".unknown .ds prompt", AiConfig())
|
||||
|
||||
self.assertTrue(invocation.use_deepseek)
|
||||
self.assertEqual(invocation.prompt, "prompt")
|
||||
|
||||
def test_rejects_missing_prompt(self) -> None:
|
||||
with self.assertRaises(AgentSyntaxError):
|
||||
parse_agent_invocation("10 .ds .th", AiConfig())
|
||||
|
||||
def test_rejects_context_count_above_configured_limit(self) -> None:
|
||||
with self.assertRaises(AgentSyntaxError):
|
||||
parse_agent_invocation("101 prompt", AiConfig(max_context_messages=100))
|
||||
|
||||
|
||||
class AgentContextTests(unittest.IsolatedAsyncioTestCase):
|
||||
async def test_reply_is_injected_once_and_history_is_chronological(self) -> None:
|
||||
older = SimpleNamespace(id=1, chat=SimpleNamespace(id=99), text="older")
|
||||
newer = SimpleNamespace(id=2, chat=SimpleNamespace(id=99), text="newer")
|
||||
command = SimpleNamespace(
|
||||
id=3,
|
||||
chat=SimpleNamespace(id=99),
|
||||
reply_to_message=older,
|
||||
)
|
||||
|
||||
async def history():
|
||||
yield newer
|
||||
yield older
|
||||
|
||||
client = SimpleNamespace(get_chat_history=lambda *args, **kwargs: history())
|
||||
invocation = parse_agent_invocation("2 prompt", AiConfig())
|
||||
with tempfile.TemporaryDirectory() as temp_dir:
|
||||
messages = await _build_context_messages(
|
||||
client,
|
||||
command,
|
||||
invocation,
|
||||
AiConfig(),
|
||||
Path(temp_dir),
|
||||
user_template="wrapped={{message.is_reply}} {{text}}",
|
||||
template_context={},
|
||||
)
|
||||
|
||||
self.assertEqual(len(messages), 2)
|
||||
self.assertIn("wrapped=true", messages[0].message())
|
||||
self.assertIn("older", messages[0].message())
|
||||
self.assertIn("newer", messages[1].message())
|
||||
|
||||
async def test_media_over_limit_is_not_downloaded(self) -> None:
|
||||
message = SimpleNamespace(
|
||||
id=1,
|
||||
document=SimpleNamespace(file_size=6, mime_type="text/plain"),
|
||||
)
|
||||
client = SimpleNamespace(download_media=AsyncMock())
|
||||
|
||||
with tempfile.TemporaryDirectory() as temp_dir:
|
||||
media = await _download_message_media(client, message, Path(temp_dir), 5)
|
||||
|
||||
self.assertEqual(media, [])
|
||||
client.download_media.assert_not_awaited()
|
||||
|
||||
async def test_media_within_limit_is_attached(self) -> None:
|
||||
with tempfile.TemporaryDirectory() as temp_dir:
|
||||
path = Path(temp_dir) / "note.txt"
|
||||
path.write_bytes(b"hello")
|
||||
message = SimpleNamespace(
|
||||
id=1,
|
||||
document=SimpleNamespace(
|
||||
file_size=5,
|
||||
mime_type="text/plain",
|
||||
file_name="note.txt",
|
||||
),
|
||||
)
|
||||
client = SimpleNamespace(download_media=AsyncMock(return_value=str(path)))
|
||||
|
||||
media = await _download_message_media(client, message, Path(temp_dir), 5)
|
||||
|
||||
self.assertEqual(len(media), 1)
|
||||
self.assertEqual(media[0].data, b"hello")
|
||||
self.assertEqual(media[0].kind, "file")
|
||||
|
||||
async def test_deepseek_and_reasoning_are_passed_to_agent(self) -> None:
|
||||
config = AiConfig(
|
||||
glm=AiProviderConfig("https://glm.example", "glm-key", "glm-model"),
|
||||
deepseek=AiProviderConfig("https://ds.example", "ds-key", "ds-model"),
|
||||
brave_api_key="brave-key",
|
||||
)
|
||||
invocation = parse_agent_invocation(".ds .th prompt", config)
|
||||
command = SimpleNamespace(
|
||||
id=3,
|
||||
chat=SimpleNamespace(id=99),
|
||||
from_user=SimpleNamespace(id=42),
|
||||
reply_to_message=None,
|
||||
)
|
||||
account = SimpleNamespace(
|
||||
me=SimpleNamespace(id=7, full_name="Owner", username="owner"),
|
||||
config=SimpleNamespace(name="personal", session="personal"),
|
||||
)
|
||||
final_message = SimpleNamespace(message=lambda: "answer")
|
||||
session = SimpleNamespace(
|
||||
run=AsyncMock(
|
||||
return_value=SimpleNamespace(
|
||||
succeeded=True,
|
||||
final_message=final_message,
|
||||
error=None,
|
||||
)
|
||||
)
|
||||
)
|
||||
agent = MagicMock()
|
||||
agent.create_session.return_value = session
|
||||
|
||||
with (
|
||||
patch("app.agentic.OpenAIBasedLLM") as llm_class,
|
||||
patch("app.agentic.Aggregator", return_value=agent) as aggregator_class,
|
||||
patch("app.agentic.BraveSearchTools") as brave_tools,
|
||||
patch("app.agentic._build_context_messages", AsyncMock(return_value=[])),
|
||||
):
|
||||
brave_tools.return_value.tools.return_value = []
|
||||
result = await run_agent_turn(
|
||||
SimpleNamespace(),
|
||||
command,
|
||||
account,
|
||||
AgentSessionStore(),
|
||||
invocation,
|
||||
config,
|
||||
)
|
||||
|
||||
self.assertEqual(result.answer, "answer")
|
||||
llm_class.assert_called_once_with(
|
||||
api_url="https://ds.example",
|
||||
api_key="ds-key",
|
||||
model="ds-model",
|
||||
)
|
||||
self.assertTrue(aggregator_class.call_args.kwargs["thinking"])
|
||||
brave_tools.assert_called_once_with(api_key="brave-key", timeout_seconds=10.0)
|
||||
result.handle.close()
|
||||
|
||||
async def test_reply_to_ai_response_reuses_existing_session(self) -> None:
|
||||
workspace = MagicMock()
|
||||
final_message = SimpleNamespace(message=lambda: "continued")
|
||||
session = SimpleNamespace(
|
||||
context={"old": True},
|
||||
run=AsyncMock(
|
||||
return_value=SimpleNamespace(
|
||||
succeeded=True,
|
||||
final_message=final_message,
|
||||
error=None,
|
||||
)
|
||||
),
|
||||
)
|
||||
handle = AgentSessionHandle(session, workspace, set())
|
||||
store = AgentSessionStore()
|
||||
ai_response = SimpleNamespace(id=50, chat=SimpleNamespace(id=99))
|
||||
store.register_response(handle, ai_response)
|
||||
command = SimpleNamespace(
|
||||
id=51,
|
||||
chat=SimpleNamespace(id=99),
|
||||
from_user=SimpleNamespace(id=7, full_name="Owner"),
|
||||
reply_to_message=ai_response,
|
||||
)
|
||||
account = SimpleNamespace(
|
||||
me=SimpleNamespace(id=7, full_name="Owner"),
|
||||
config=SimpleNamespace(name="personal", session="personal"),
|
||||
)
|
||||
|
||||
result = await run_agent_turn(
|
||||
SimpleNamespace(),
|
||||
command,
|
||||
account,
|
||||
store,
|
||||
parse_agent_invocation("continue", AiConfig()),
|
||||
AiConfig(),
|
||||
)
|
||||
|
||||
self.assertIs(result.handle, handle)
|
||||
self.assertEqual(result.answer, "continued")
|
||||
session.run.assert_awaited_once_with("continue")
|
||||
self.assertTrue(session.context["is_owner"])
|
||||
store.clear()
|
||||
workspace.cleanup.assert_called_once()
|
||||
|
||||
|
||||
class AiCommandTests(unittest.IsolatedAsyncioTestCase):
|
||||
def _ctx(self, args: str = "prompt") -> SimpleNamespace:
|
||||
sent = SimpleNamespace(id=124, chat=SimpleNamespace(id=99))
|
||||
message = SimpleNamespace(
|
||||
id=123,
|
||||
react=AsyncMock(),
|
||||
reply_text=AsyncMock(return_value=sent),
|
||||
)
|
||||
session_store = MagicMock()
|
||||
return SimpleNamespace(
|
||||
args=args,
|
||||
message=message,
|
||||
client=SimpleNamespace(),
|
||||
account=SimpleNamespace(
|
||||
me=SimpleNamespace(id=7, full_name="Owner"),
|
||||
config=SimpleNamespace(name="personal", session="personal"),
|
||||
ai_sessions=session_store,
|
||||
),
|
||||
manager=SimpleNamespace(config=SimpleNamespace(ai=AiConfig())),
|
||||
reply=message.reply_text,
|
||||
)
|
||||
|
||||
def test_ai_is_registered_as_shared(self) -> None:
|
||||
definition = command_registry.resolve("ai")
|
||||
|
||||
self.assertIsNotNone(definition)
|
||||
self.assertTrue(definition.is_shared)
|
||||
self.assertEqual(definition.guest_prefix, "..")
|
||||
|
||||
async def test_success_replies_with_agent_answer(self) -> None:
|
||||
ctx = self._ctx()
|
||||
handle = MagicMock()
|
||||
with patch(
|
||||
"app.commands.builtins.run_agent_turn",
|
||||
AsyncMock(return_value=SimpleNamespace(answer="answer", handle=handle)),
|
||||
):
|
||||
await ai(ctx)
|
||||
|
||||
ctx.reply.assert_awaited_once_with("answer")
|
||||
ctx.account.ai_sessions.register_response.assert_called_once_with(
|
||||
handle,
|
||||
ctx.reply.return_value,
|
||||
)
|
||||
ctx.message.react.assert_not_awaited()
|
||||
|
||||
async def test_invalid_syntax_only_dislikes(self) -> None:
|
||||
ctx = self._ctx(".ds .th")
|
||||
|
||||
with self.assertLogs("app.commands.builtins", level="ERROR"):
|
||||
await ai(ctx)
|
||||
|
||||
ctx.message.react.assert_awaited_once_with("👎")
|
||||
ctx.reply.assert_not_awaited()
|
||||
|
||||
async def test_agent_error_only_dislikes(self) -> None:
|
||||
ctx = self._ctx()
|
||||
with patch(
|
||||
"app.commands.builtins.run_agent_turn",
|
||||
AsyncMock(side_effect=RuntimeError("API failed")),
|
||||
):
|
||||
with self.assertLogs("app.commands.builtins", level="ERROR"):
|
||||
await ai(ctx)
|
||||
|
||||
ctx.message.react.assert_awaited_once_with("👎")
|
||||
ctx.reply.assert_not_awaited()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user