1197 lines
46 KiB
Python
1197 lines
46 KiB
Python
import asyncio
|
||
import base64
|
||
import contextlib
|
||
import functools
|
||
import tempfile
|
||
from pathlib import Path
|
||
from types import SimpleNamespace
|
||
from typing import Any, cast
|
||
|
||
import pytest
|
||
from aiogram.exceptions import TelegramBadRequest, TelegramNetworkError
|
||
from aiogram.methods import SendMessage
|
||
from aiogram.types import Update
|
||
from claude_agent_sdk import PermissionResultAllow, PermissionResultDeny
|
||
|
||
from beaver_gateway.app import McpRegistry
|
||
from beaver_gateway.backends.transcript import build_entries
|
||
from beaver_gateway.conversations.tools import _tools
|
||
from beaver_gateway.frontends.base import GatewayRuntime
|
||
from beaver_gateway.frontends.files import SendFileError, SendFiles, method_for
|
||
from beaver_gateway.frontends.telegram import Attachments, TelegramFrontend
|
||
from beaver_gateway.frontends.telegram.drafts import Draft
|
||
from beaver_gateway.frontends.telegram.render import (
|
||
LIMIT,
|
||
chunks,
|
||
status_label,
|
||
to_html,
|
||
to_html_tail,
|
||
)
|
||
from beaver_gateway.storage.models import ConversationBinding, Delivery, TelegramUpdate
|
||
from sqlmodel import select
|
||
from test_conversations import ScriptedClient, StubFrontend, World
|
||
|
||
USER = 42
|
||
GENERAL = f"{USER}/901"
|
||
|
||
|
||
class FakeBot:
|
||
def __init__(self) -> None:
|
||
self.updates: list[dict[str, Any]] = []
|
||
self.sent: list[dict[str, Any]] = []
|
||
self.files: list[dict[str, Any]] = []
|
||
self.reject_files = False
|
||
self.drafts: list[dict[str, Any]] = []
|
||
self.edits: list[dict[str, Any]] = []
|
||
self.topics: list[str] = []
|
||
self.reactions: list[tuple[int, list[Any]]] = []
|
||
self.fail_sends = 0
|
||
self.reject_html = False
|
||
self.gone_threads: set[int] = set()
|
||
self.order: list[tuple[str, int | None]] = []
|
||
self._message_id = 100
|
||
self.session = SimpleNamespace(close=self._close)
|
||
|
||
async def _close(self) -> None:
|
||
pass
|
||
|
||
async def get_me(self) -> Any:
|
||
return SimpleNamespace(username="bot", has_topics_enabled=True)
|
||
|
||
async def get_updates(self, offset=None, **_: Any) -> list[Update]:
|
||
await asyncio.sleep(0.01)
|
||
return [
|
||
Update.model_validate(u)
|
||
for u in self.updates
|
||
if offset is None or u["update_id"] >= offset
|
||
]
|
||
|
||
async def send_message(
|
||
self, chat_id, text, message_thread_id=None, parse_mode=None, **kwargs: Any
|
||
) -> Any:
|
||
if self.fail_sends:
|
||
self.fail_sends -= 1
|
||
raise TelegramNetworkError(
|
||
method=SendMessage(chat_id=0, text="x"), message="boom"
|
||
)
|
||
if self.reject_html and parse_mode == "HTML":
|
||
raise TelegramBadRequest(
|
||
method=SendMessage(chat_id=0, text="x"),
|
||
message="Bad Request: can't parse entities",
|
||
)
|
||
if message_thread_id in self.gone_threads:
|
||
raise TelegramBadRequest(
|
||
method=SendMessage(chat_id=0, text="x"),
|
||
message="Bad Request: message thread not found",
|
||
)
|
||
self._message_id += 1
|
||
self.order.append(("message", message_thread_id))
|
||
self.sent.append(
|
||
{
|
||
"chat_id": chat_id,
|
||
"thread": message_thread_id,
|
||
"text": text,
|
||
"parse_mode": parse_mode,
|
||
"markup": kwargs.get("reply_markup"),
|
||
"message_id": self._message_id,
|
||
}
|
||
)
|
||
return SimpleNamespace(message_id=self._message_id)
|
||
|
||
async def send_message_draft(self, **kwargs: Any) -> bool:
|
||
await asyncio.sleep(0.08)
|
||
self.drafts.append(kwargs)
|
||
self.order.append(("draft", kwargs.get("message_thread_id")))
|
||
return True
|
||
|
||
async def create_forum_topic(self, chat_id, name, **_: Any) -> Any:
|
||
self.topics.append(name)
|
||
created = sum(1 for t in self.topics if not t.startswith("edit:"))
|
||
return SimpleNamespace(message_thread_id=900 + created, name=name)
|
||
|
||
async def edit_forum_topic(
|
||
self, chat_id, message_thread_id, name=None, **_: Any
|
||
) -> bool:
|
||
self.topics.append(f"edit:{message_thread_id}:{name}")
|
||
return True
|
||
|
||
async def set_message_reaction(
|
||
self, chat_id, message_id, reaction=None, **_: Any
|
||
) -> bool:
|
||
self.reactions.append((message_id, reaction or []))
|
||
return True
|
||
|
||
async def answer_callback_query(self, *_: Any, **__: Any) -> bool:
|
||
return True
|
||
|
||
async def edit_message_text(self, text, chat_id, message_id, **_: Any) -> Any:
|
||
self.edits.append({"message_id": message_id, "text": text})
|
||
return True
|
||
|
||
async def edit_message_reply_markup(
|
||
self, chat_id, message_id, reply_markup=None
|
||
) -> Any:
|
||
self.edits.append({"message_id": message_id, "markup": reply_markup})
|
||
return True
|
||
|
||
async def download(self, file_id, destination=None) -> None:
|
||
Path(destination).write_bytes(b"data")
|
||
|
||
async def _upload(self, method: str, chat_id: int, **kwargs: Any) -> Any:
|
||
if self.reject_files:
|
||
raise TelegramBadRequest(
|
||
method=SendMessage(chat_id=0, text="x"),
|
||
message="Bad Request: PHOTO_INVALID_DIMENSIONS",
|
||
)
|
||
self._message_id += 1
|
||
self.files.append({"method": method, "chat_id": chat_id, **kwargs})
|
||
return SimpleNamespace(message_id=self._message_id)
|
||
|
||
def __getattr__(self, name: str) -> Any:
|
||
if name.startswith("send_"):
|
||
return functools.partial(self._upload, name.removeprefix("send_"))
|
||
raise AttributeError(name)
|
||
|
||
# helpers for tests
|
||
def push(self, payload: dict[str, Any]) -> None:
|
||
payload["update_id"] = 1000 + len(self.updates)
|
||
self.updates.append(payload)
|
||
|
||
def message(self, text: str, *, thread: int | None = None, uid: int = USER) -> None:
|
||
body: dict[str, Any] = {
|
||
"message_id": 10 + len(self.updates),
|
||
"date": 1700000000,
|
||
"chat": {"id": USER, "type": "private"},
|
||
"from": {"id": uid, "is_bot": False, "first_name": "h"},
|
||
"text": text,
|
||
}
|
||
if thread is not None:
|
||
body |= {"message_thread_id": thread, "is_topic_message": True}
|
||
self.push({"message": body})
|
||
|
||
def topic_created(self, name: str, thread: int) -> None:
|
||
self.push(
|
||
{
|
||
"message": {
|
||
"message_id": 10 + len(self.updates),
|
||
"date": 1700000000,
|
||
"chat": {"id": USER, "type": "private"},
|
||
"from": {"id": USER, "is_bot": False, "first_name": "h"},
|
||
"message_thread_id": thread,
|
||
"is_topic_message": True,
|
||
"forum_topic_created": {"name": name, "icon_color": 1},
|
||
}
|
||
}
|
||
)
|
||
|
||
def topic_edited(self, name: str, thread: int) -> None:
|
||
self.push(
|
||
{
|
||
"message": {
|
||
"message_id": 10 + len(self.updates),
|
||
"date": 1700000000,
|
||
"chat": {"id": USER, "type": "private"},
|
||
"from": {"id": USER, "is_bot": False, "first_name": "h"},
|
||
"message_thread_id": thread,
|
||
"is_topic_message": True,
|
||
"forum_topic_edited": {"name": name},
|
||
}
|
||
}
|
||
)
|
||
|
||
def media(self, body: dict[str, Any], *, thread: int | None = None) -> None:
|
||
message: dict[str, Any] = {
|
||
"message_id": 10 + len(self.updates),
|
||
"date": 1700000000,
|
||
"chat": {"id": USER, "type": "private"},
|
||
"from": {"id": USER, "is_bot": False, "first_name": "h"},
|
||
**body,
|
||
}
|
||
if thread is not None:
|
||
message |= {"message_thread_id": thread, "is_topic_message": True}
|
||
self.push({"message": message})
|
||
|
||
def callback(self, data: str, message_id: int) -> None:
|
||
self.push(
|
||
{
|
||
"callback_query": {
|
||
"id": f"cb{len(self.updates)}",
|
||
"from": {"id": USER, "is_bot": False, "first_name": "h"},
|
||
"chat_instance": "ci",
|
||
"data": data,
|
||
"message": {
|
||
"message_id": message_id,
|
||
"date": 1700000000,
|
||
"chat": {"id": USER, "type": "private"},
|
||
"text": "?",
|
||
},
|
||
}
|
||
}
|
||
)
|
||
|
||
|
||
class Stack:
|
||
def __init__(self, world: World, bot: FakeBot) -> None:
|
||
self.world = world
|
||
self.bot = bot
|
||
self.tg = TelegramFrontend(
|
||
token="t",
|
||
user_id=USER,
|
||
master_agent="a",
|
||
branch_agent="a",
|
||
draft_interval=0.05,
|
||
outbox_backoff=0.01,
|
||
)
|
||
self.tg._bot = bot # noqa: SLF001
|
||
markdown = StubFrontend("markdown", ("deep",), {"deep": "d"}, home=True)
|
||
markdown.conversations = world.conversations
|
||
frontends = [self.tg, markdown]
|
||
world.conversations._frontends = frontends # noqa: SLF001
|
||
runtime = GatewayRuntime(
|
||
agents=world.conversations._agents, # noqa: SLF001
|
||
mcps=McpRegistry([]),
|
||
backends=world.conversations._backends, # noqa: SLF001
|
||
token_store=None,
|
||
db=world.db,
|
||
conversations=world.conversations,
|
||
bus=world.bus,
|
||
pool=world.pool,
|
||
frontends=tuple(frontends),
|
||
)
|
||
self.tg.configure(runtime)
|
||
self.tasks = [
|
||
asyncio.create_task(self.tg.inbox.run()),
|
||
asyncio.create_task(self.tg.outbox.run()),
|
||
asyncio.create_task(self.tg._events()), # noqa: SLF001
|
||
]
|
||
|
||
async def close(self) -> None:
|
||
for t in self.tasks:
|
||
t.cancel()
|
||
for t in self.tasks:
|
||
with pytest.raises(BaseException):
|
||
await t
|
||
|
||
async def until(self, pred, timeout: float = 5.0, what: str = "condition") -> Any:
|
||
deadline = asyncio.get_running_loop().time() + timeout
|
||
while asyncio.get_running_loop().time() < deadline:
|
||
value = pred()
|
||
if value:
|
||
return value
|
||
await asyncio.sleep(0.02)
|
||
msg = f"{what} never happened; sent={self.bot.sent}"
|
||
raise AssertionError(msg)
|
||
|
||
def sent_with(self, needle: str) -> dict[str, Any] | None:
|
||
return next((m for m in self.bot.sent if needle in m["text"]), None)
|
||
|
||
|
||
@pytest.fixture
|
||
async def stack() -> Stack:
|
||
root = Path(tempfile.mkdtemp(prefix="beaver-tg-"))
|
||
world = await World(root).setup()
|
||
await world.conversations.start()
|
||
s = Stack(world, FakeBot())
|
||
yield s
|
||
await s.close()
|
||
await world.conversations.stop()
|
||
await world.pool.close_all()
|
||
await world.db.dispose()
|
||
|
||
|
||
async def test_general_is_master_and_reply_has_no_thread(stack: Stack) -> None:
|
||
stack.bot.message("hi")
|
||
reply = await stack.until(lambda: stack.sent_with("hi"), what="reply")
|
||
assert reply["text"].startswith("ok:[seed: clean] master")
|
||
assert reply["thread"] == 901
|
||
assert reply["parse_mode"] == "HTML"
|
||
assert stack.bot.topics == ["🦫 General"]
|
||
master = await stack.world.conversations.find_bound(
|
||
frontend="telegram", external_id=GENERAL
|
||
)
|
||
assert master is not None and master.kind == "master"
|
||
assert stack.bot.drafts and stack.bot.drafts[0]["chat_id"] == USER
|
||
await asyncio.sleep(0.2)
|
||
final = max(i for i, o in enumerate(stack.bot.order) if o[0] == "message")
|
||
assert all(o[0] != "draft" for o in stack.bot.order[final:])
|
||
last = stack.bot.drafts[-1]
|
||
assert last["text"] == reply["text"] and last["parse_mode"] == "HTML"
|
||
assert stack.bot.drafts[0]["text"] == "⏳ thinking"
|
||
rows = await stack.world.conversations.queue.recent(master.id)
|
||
assert [r.origin for r in rows] == ["telegram"]
|
||
|
||
|
||
async def test_new_topic_becomes_morning_branch_and_replies_in_thread(
|
||
stack: Stack,
|
||
) -> None:
|
||
stack.bot.topic_created("план", 7)
|
||
stack.bot.message("hello topic", thread=7)
|
||
reply = await stack.until(lambda: stack.sent_with("hello topic"), what="reply")
|
||
assert reply["thread"] == 7
|
||
assert len([m for m in stack.bot.sent if m["thread"] == 7]) == 1
|
||
branch = await stack.world.conversations.find_bound(
|
||
frontend="telegram", external_id=f"{USER}/7"
|
||
)
|
||
assert branch is not None
|
||
assert branch.kind == "branch" and branch.title == "план"
|
||
master = await stack.world.conversations.find_bound(
|
||
frontend="telegram", external_id=GENERAL
|
||
)
|
||
assert branch.parent_id == master.id
|
||
prompts = [p for c in ScriptedClient.instances for p in c.prompts]
|
||
seed = next(p for p in prompts if p.startswith("[seed: morning] branch «план»"))
|
||
assert "No handout arrived." in seed
|
||
assert seed.endswith("hello topic")
|
||
assert stack.bot.drafts[-1]["message_thread_id"] == 7
|
||
|
||
|
||
async def test_first_message_without_service_message_seeds_with_text(
|
||
stack: Stack,
|
||
) -> None:
|
||
stack.bot.message("сразу текстом", thread=8)
|
||
await stack.until(lambda: stack.sent_with("сразу текстом"), what="reply")
|
||
branch = await stack.world.conversations.find_bound(
|
||
frontend="telegram", external_id=f"{USER}/8"
|
||
)
|
||
assert branch.title == "сразу текстом"
|
||
client = next(
|
||
c for c in ScriptedClient.instances if c.prompts and "сразу" in c.prompts[0]
|
||
)
|
||
assert client.prompts[0].startswith("[seed: morning]")
|
||
assert client.prompts[0].endswith("сразу текстом")
|
||
assert len(client.prompts) == 1
|
||
|
||
|
||
async def test_message_into_merged_branch_rebinds_the_topic(stack: Stack) -> None:
|
||
stack.bot.message("one", thread=9)
|
||
await stack.until(lambda: stack.sent_with("ok:"), what="first reply")
|
||
old = await stack.world.conversations.find_bound(
|
||
frontend="telegram", external_id=f"{USER}/9"
|
||
)
|
||
await stack.world.conversations.set_status(old, "merged")
|
||
stack.bot.message("two", thread=9)
|
||
await stack.until(lambda: stack.sent_with("two"), what="second reply")
|
||
new = await stack.world.conversations.find_bound(
|
||
frontend="telegram", external_id=f"{USER}/9"
|
||
)
|
||
assert new.id != old.id and new.status == "open"
|
||
async with stack.world.db.session() as session:
|
||
rows = list(
|
||
(
|
||
await session.exec(
|
||
select(ConversationBinding).where(
|
||
ConversationBinding.external_id == f"{USER}/9"
|
||
)
|
||
)
|
||
).all()
|
||
)
|
||
assert {(r.conversation_id, r.visible) for r in rows} == {
|
||
(old.id, False),
|
||
(new.id, True),
|
||
}
|
||
|
||
|
||
async def test_reply_from_another_window_is_mirrored_with_marker(stack: Stack) -> None:
|
||
stack.bot.message("hi")
|
||
await stack.until(lambda: stack.sent_with("hi"), what="reply")
|
||
master = await stack.world.conversations.find_bound(
|
||
frontend="telegram", external_id=GENERAL
|
||
)
|
||
await stack.world.conversations.post(master, "from panel", origin="user")
|
||
await stack.until(lambda: stack.sent_with("ok:from panel"), what="mirrored reply")
|
||
texts = [m["text"] for m in stack.bot.sent]
|
||
marker = next(i for i, t in enumerate(texts) if t.startswith("📝 from the panel:"))
|
||
assert "from panel" in texts[marker]
|
||
assert texts[marker + 1] == "ok:from panel"
|
||
|
||
|
||
async def test_say_is_delivered_and_inject_turns_are_silent(stack: Stack) -> None:
|
||
stack.bot.message("hi")
|
||
await stack.until(lambda: stack.sent_with("hi"), what="reply")
|
||
master = await stack.world.conversations.find_bound(
|
||
frontend="telegram", external_id=GENERAL
|
||
)
|
||
before = len(stack.bot.sent)
|
||
await stack.world.conversations.inject(
|
||
master, "cron tick", urgency="urgent", origin="крон"
|
||
)
|
||
await stack.world.settle(master, 2)
|
||
await asyncio.sleep(0.2)
|
||
assert len(stack.bot.sent) == before
|
||
await stack.world.conversations.say(master, "psst")
|
||
await stack.until(lambda: stack.sent_with("psst"), what="say")
|
||
|
||
|
||
async def test_question_becomes_buttons_and_callback_answers(stack: Stack) -> None:
|
||
stack.bot.message("hi")
|
||
await stack.until(lambda: stack.sent_with("hi"), what="reply")
|
||
master = await stack.world.conversations.find_bound(
|
||
frontend="telegram", external_id=GENERAL
|
||
)
|
||
payload = {
|
||
"questions": [
|
||
{
|
||
"header": "Цвет",
|
||
"question": "Какой цвет?",
|
||
"options": [
|
||
{"label": "Синий", "description": "как небо"},
|
||
{"label": "Красный", "description": ""},
|
||
],
|
||
"multiSelect": False,
|
||
}
|
||
]
|
||
}
|
||
asking = asyncio.create_task(
|
||
stack.world.conversations.ask(master.external_id, payload)
|
||
)
|
||
question = await stack.until(
|
||
lambda: stack.sent_with("Какой цвет?"), what="question"
|
||
)
|
||
assert question["markup"] is not None
|
||
labels = [b.text for row in question["markup"].inline_keyboard for b in row]
|
||
assert labels == ["Синий", "Красный"]
|
||
pending = stack.world.conversations.pending_question(master.external_id)
|
||
assert pending is not None
|
||
stack.bot.callback(f"q:{pending[0]}:0:1", question["message_id"])
|
||
assert await asyncio.wait_for(asking, 5) == "Красный"
|
||
assert stack.world.conversations.answer_text("Красный") == (
|
||
"The user answered: Красный"
|
||
)
|
||
await stack.until(
|
||
lambda: any("✅ Красный" in e.get("text", "") for e in stack.bot.edits),
|
||
what="edit",
|
||
)
|
||
row = await stack.world.conversations.get(master.external_id)
|
||
assert row.pending_question is False
|
||
|
||
|
||
async def test_question_timeout_renders_text_and_free_text_answers(
|
||
stack: Stack,
|
||
) -> None:
|
||
stack.bot.message("hi")
|
||
await stack.until(lambda: stack.sent_with("hi"), what="reply")
|
||
master = await stack.world.conversations.find_bound(
|
||
frontend="telegram", external_id=GENERAL
|
||
)
|
||
stack.world.conversations._question_timeout = 0.3 # noqa: SLF001
|
||
payload = {"questions": [{"header": "Q", "question": "Сколько?", "options": []}]}
|
||
result = await stack.world.conversations.ask(master.external_id, payload)
|
||
assert result is None
|
||
assert "did not answer" in stack.world.conversations.answer_text(None)
|
||
await stack.until(
|
||
lambda: any("time is up" in e.get("text", "") for e in stack.bot.edits),
|
||
what="timeout edit",
|
||
)
|
||
stack.world.conversations._question_timeout = 5.0 # noqa: SLF001
|
||
asking = asyncio.create_task(
|
||
stack.world.conversations.ask(master.external_id, payload)
|
||
)
|
||
await stack.until(
|
||
lambda: stack.world.conversations.pending_question(master.external_id),
|
||
what="pending",
|
||
)
|
||
stack.bot.message("семь")
|
||
assert await asyncio.wait_for(asking, 5) == "семь"
|
||
|
||
|
||
async def test_commands_status_merge_and_new(stack: Stack) -> None:
|
||
stack.bot.message("/status")
|
||
status = await stack.until(lambda: stack.sent_with("master · open"), what="status")
|
||
assert "pool:" in status["text"]
|
||
stack.bot.message("work", thread=11)
|
||
await stack.until(lambda: stack.sent_with("work"), what="branch reply")
|
||
branch = await stack.world.conversations.find_bound(
|
||
frontend="telegram", external_id=f"{USER}/11"
|
||
)
|
||
await stack.world.store.append(
|
||
stack.world.key(branch.session_id),
|
||
build_entries(
|
||
[{"role": "user", "content": "q"}, {"role": "assistant", "content": "a"}],
|
||
session_id=branch.session_id,
|
||
cwd=str(stack.world.root),
|
||
model="m",
|
||
),
|
||
)
|
||
stack.bot.message("/merge", thread=11)
|
||
await stack.until(lambda: stack.sent_with("merged into the master"), what="merged")
|
||
branch = await stack.world.conversations.find_bound(
|
||
frontend="telegram", external_id=f"{USER}/11"
|
||
)
|
||
assert branch.status == "merged"
|
||
assert stack.bot.topics == ["🦫 General", "edit:11:✅ work"]
|
||
stack.bot.message("/new отчёт")
|
||
await stack.until(lambda: stack.sent_with("in a new topic"), what="new topic")
|
||
assert stack.bot.topics == ["🦫 General", "edit:11:✅ work", "отчёт"]
|
||
child = await stack.world.conversations.find_bound(
|
||
frontend="telegram", external_id=f"{USER}/902"
|
||
)
|
||
assert child is not None and child.title == "отчёт"
|
||
await asyncio.sleep(0.2)
|
||
assert await stack.world.conversations.queue.recent(child.id) == []
|
||
assert all(m["thread"] != 902 for m in stack.bot.sent)
|
||
assert all(d["message_thread_id"] != 902 for d in stack.bot.drafts)
|
||
stack.bot.message("первое в новый топик", thread=902)
|
||
reply = await stack.until(lambda: stack.sent_with("первое в новый топик"), what="r")
|
||
assert reply["thread"] == 902
|
||
assert reply["text"].startswith("ok:[seed: morning] branch «отчёт»")
|
||
assert await stack.tg.mark_topic(child)
|
||
assert stack.bot.topics[-1] == "edit:902:✅ отчёт"
|
||
|
||
|
||
async def test_brief_branch_from_the_tool_replies_in_its_topic(stack: Stack) -> None:
|
||
stack.bot.message("hi")
|
||
await stack.until(lambda: stack.sent_with("hi"), what="reply")
|
||
master = await stack.world.conversations.find_bound(
|
||
frontend="telegram", external_id=GENERAL
|
||
)
|
||
child = await stack.world.conversations.spawn(
|
||
kind="branch",
|
||
seed="brief",
|
||
text="найди X",
|
||
title="X",
|
||
parent=master,
|
||
origin="mcp",
|
||
)
|
||
reply = await stack.until(lambda: stack.sent_with("найди X"), what="brief reply")
|
||
assert reply["thread"] == 902
|
||
assert not any(m["text"].startswith("📝") for m in stack.bot.sent)
|
||
assert any(d["message_thread_id"] == 902 for d in stack.bot.drafts)
|
||
assert (await stack.world.conversations.get(child.external_id)).flags == {}
|
||
|
||
|
||
async def test_master_topic_survives_rotation(stack: Stack) -> None:
|
||
stack.bot.message("hi")
|
||
await stack.until(lambda: stack.sent_with("hi"), what="reply")
|
||
master = await stack.world.conversations.find_bound(
|
||
frontend="telegram", external_id=GENERAL
|
||
)
|
||
stack.bot.message("into general", thread=901)
|
||
await stack.until(lambda: stack.sent_with("into general"), what="reply")
|
||
assert (await stack.world.conversations.find(kind="branch")) == []
|
||
await stack.world.conversations.set_status(master, "closed")
|
||
stack.bot.message("after rotation", thread=901)
|
||
await stack.until(lambda: stack.sent_with("after rotation"), what="reply")
|
||
fresh = await stack.world.conversations.find_bound(
|
||
frontend="telegram", external_id=GENERAL
|
||
)
|
||
assert fresh.kind == "master" and fresh.id != master.id
|
||
assert stack.bot.topics == ["🦫 General"]
|
||
assert (await stack.world.conversations.find(kind="branch")) == []
|
||
|
||
|
||
async def test_inbox_stores_first_and_replays_after_restart(stack: Stack) -> None:
|
||
stack.bot.message("hi")
|
||
await stack.until(lambda: stack.sent_with("hi"), what="reply")
|
||
async with stack.world.db.session() as session:
|
||
rows = list((await session.exec(select(TelegramUpdate))).all())
|
||
assert [r.update_id for r in rows] == [1000]
|
||
assert rows[0].processed_at is not None and rows[0].error is None
|
||
async with stack.world.db.session() as session:
|
||
session.add(
|
||
TelegramUpdate(
|
||
update_id=999,
|
||
payload={
|
||
"update_id": 999,
|
||
"message": {
|
||
"message_id": 1,
|
||
"date": 1700000000,
|
||
"chat": {"id": USER, "type": "private"},
|
||
"from": {"id": USER, "is_bot": False, "first_name": "h"},
|
||
"text": "replayed",
|
||
},
|
||
},
|
||
)
|
||
)
|
||
await session.commit()
|
||
await stack.until(
|
||
lambda: stack.sent_with("replayed"), timeout=8, what="replayed reply"
|
||
)
|
||
stack.bot.message("stranger", uid=1)
|
||
await asyncio.sleep(0.3)
|
||
assert stack.sent_with("stranger") is None
|
||
|
||
|
||
async def test_outbox_retries_and_falls_back_to_plain(stack: Stack) -> None:
|
||
stack.bot.fail_sends = 2
|
||
stack.bot.reject_html = True
|
||
stack.bot.message("**bold**")
|
||
reply = await stack.until(lambda: stack.sent_with("**bold**"), what="reply")
|
||
assert reply["parse_mode"] is None
|
||
async with stack.world.db.session() as session:
|
||
rows = list((await session.exec(select(Delivery))).all())
|
||
sent = [r for r in rows if r.status == "sent"]
|
||
assert sent and sent[0].plain and sent[0].attempts >= 3
|
||
|
||
|
||
async def test_can_use_tool_answers_through_deny(stack: Stack) -> None:
|
||
seen: list[tuple[str, dict[str, Any]]] = []
|
||
|
||
async def asker(key: str, payload: dict[str, Any]) -> str:
|
||
seen.append((key, payload))
|
||
return "Пользователь ответил: да"
|
||
|
||
backend = stack.world.backend
|
||
backend._asker = asker # noqa: SLF001
|
||
can_use_tool = backend._can_use_tool("conv-1") # noqa: SLF001
|
||
allow = await can_use_tool("Bash", {"command": "ls"}, None)
|
||
assert isinstance(allow, PermissionResultAllow)
|
||
deny = await can_use_tool("AskUserQuestion", {"questions": []}, None)
|
||
assert isinstance(deny, PermissionResultDeny)
|
||
assert deny.message == "Пользователь ответил: да"
|
||
assert seen == [("conv-1", {"questions": []})]
|
||
|
||
|
||
def test_render_helpers() -> None:
|
||
assert (
|
||
to_html("**жирно** и `code` <b>")
|
||
== "<b>жирно</b> и <code>code</code> <b>"
|
||
)
|
||
assert to_html("# Заголовок\n- пункт") == "<b>Заголовок</b>\n• пункт"
|
||
assert to_html("```\nx < 1 & y\n```") == "<pre>x < 1 & y</pre>"
|
||
assert (
|
||
to_html("```py\nunterminated\n")
|
||
== '<pre><code class="language-py">unterminated</code></pre>'
|
||
)
|
||
assert (
|
||
to_html("[док](https://a.b/c?x=1&y=2)")
|
||
== '<a href="https://a.b/c?x=1&y=2">док</a>'
|
||
)
|
||
assert (
|
||
to_html("[tg](tg://user?id=1) https://x.y/a_b_c")
|
||
== '<a href="tg://user?id=1">tg</a> https://x.y/a_b_c'
|
||
)
|
||
assert to_html(" ") == (
|
||
'<a href="https://x.y/i.png">alt</a> <a href="https://x.y/j.png">https://x.y/j.png</a>'
|
||
)
|
||
assert to_html("~~нет~~ _к_ *к* __ж__") == "<s>нет</s> <i>к</i> <i>к</i> <b>ж</b>"
|
||
assert to_html("---\n***\n___") == "———\n———\n———"
|
||
assert to_html("`**raw** [[x]]`") == "<code>**raw** [[x]]</code>"
|
||
parts = chunks("абв\n\n" + "г" * 30 + "\n\n" + "д" * 30, limit=40)
|
||
assert parts == ["абв\n\n" + "г" * 30, "д" * 30]
|
||
assert chunks(" \n ") == []
|
||
assert status_label("Read", {"file_path": "/x"}) == "reading"
|
||
assert status_label("mcp__firefly__list_accounts", None) == "firefly: list_accounts"
|
||
assert status_label("Foo", {"command": "ls -la"}) == "Foo ls -la"
|
||
|
||
|
||
def test_render_wikilinks() -> None:
|
||
assert to_html("[[Note]]") == "<u>Note</u>"
|
||
assert to_html("[[Note|alias]]") == "<u>alias</u>"
|
||
assert to_html("[[Note#heading]]") == "<u>Note › heading</u>"
|
||
assert to_html("[[A & B|**x**]]") == "<u>**x**</u>"
|
||
|
||
|
||
def test_render_blocks() -> None:
|
||
assert (
|
||
to_html("> цитата **жирно**\n> вторая\n\nтекст")
|
||
== "<blockquote>цитата <b>жирно</b>\nвторая</blockquote>\n\nтекст"
|
||
)
|
||
assert (
|
||
to_html("1. один\n2) два\n - [ ] нет\n- [x] да")
|
||
== "1. один\n2. два\n ☐ нет\n☑ да"
|
||
)
|
||
assert (
|
||
to_html("| a | bb |\n|---|:--:|\n| ccc | d |\n| e |")
|
||
== "<pre>a | bb\nccc | d\ne</pre>"
|
||
)
|
||
assert (
|
||
to_html("```python\nprint(1)\n```")
|
||
== '<pre><code class="language-python">print(1)</code></pre>'
|
||
)
|
||
assert (
|
||
to_html("||секрет|| и `||нет||`")
|
||
== "<tg-spoiler>секрет</tg-spoiler> и <code>||нет||</code>"
|
||
)
|
||
|
||
|
||
def test_chunks_split_fence_into_valid_fences() -> None:
|
||
body = "\n".join(f"x = {i} & {i}" for i in range(400))
|
||
text = f"до\n```python\n{body}\n```\nпосле"
|
||
parts = chunks(text, limit=1500)
|
||
assert len(parts) > 2
|
||
assert parts[0] == "до"
|
||
assert all(p.startswith("```python\n") and p.endswith("```") for p in parts[1:-1])
|
||
assert parts[-1].endswith("```\nпосле")
|
||
fences = [p.removesuffix("\nпосле") for p in parts if p.startswith("```")]
|
||
inner = [p.removeprefix("```python\n").removesuffix("\n```") for p in fences]
|
||
assert "\n".join(inner) == body
|
||
for p in parts:
|
||
rendered = to_html(p)
|
||
assert 0 < len(rendered) <= 1500
|
||
for p in fences:
|
||
assert to_html(p).startswith('<pre><code class="language-python">')
|
||
assert to_html(p).endswith("</code></pre>")
|
||
|
||
|
||
def test_chunks_split_blockquote_and_stay_under_limit() -> None:
|
||
text = "\n".join(f"> строка {i} **ж** <" for i in range(300))
|
||
parts = chunks(text, limit=1000)
|
||
assert len(parts) > 1
|
||
for p in parts:
|
||
rendered = to_html(p)
|
||
assert len(rendered) <= 1000
|
||
assert rendered.startswith("<blockquote>") and rendered.endswith(
|
||
"</blockquote>"
|
||
)
|
||
assert "".join(parts).count("строка") == 300
|
||
long_line = "слово " * 2000
|
||
parts = chunks(long_line, limit=LIMIT)
|
||
assert all(0 < len(to_html(p)) <= LIMIT for p in parts)
|
||
assert " ".join(parts).split() == long_line.split()
|
||
|
||
|
||
def test_draft_breaks_a_paragraph_after_a_tool_call() -> None:
|
||
draft = Draft(cast("Any", None), chat_id=1, thread_id=None, turn_id="t")
|
||
draft.append("Checking the balance.")
|
||
draft.set_status("⏳ Bash")
|
||
draft.set_status("✍️ writing")
|
||
draft.append("It is 20 378.")
|
||
assert draft.text == "Checking the balance.\n\nIt is 20 378."
|
||
draft.set_status("⏳ Read")
|
||
draft.append("\n\nAnd a list:")
|
||
assert draft.text.endswith("It is 20 378.\n\nAnd a list:")
|
||
fresh = Draft(cast("Any", None), chat_id=1, thread_id=None, turn_id="u")
|
||
fresh.set_status("⏳ Bash")
|
||
fresh.append("First words")
|
||
assert fresh.text == "First words"
|
||
|
||
|
||
def test_to_html_tail_inside_fence() -> None:
|
||
full = (
|
||
"intro\n```py\n" + "\n".join(f"code {i} <" for i in range(200)) + "\n```\nafter"
|
||
)
|
||
rendered = to_html_tail(full, 100)
|
||
assert rendered.startswith('<pre><code class="language-py">code ')
|
||
assert rendered.endswith("</code></pre>\nafter")
|
||
assert "```" not in rendered
|
||
assert to_html_tail("a\n> b\n> c", 3) == "<blockquote>c</blockquote>"
|
||
assert to_html_tail("**x** y", 100) == "<b>x</b> y"
|
||
assert to_html_tail("x" * 50, 10) == "x" * 10
|
||
|
||
|
||
async def test_message_with_a_link_preview_is_stored_and_answered(stack: Stack) -> None:
|
||
stack.bot.push(
|
||
{
|
||
"message": {
|
||
"message_id": 77,
|
||
"date": 1700000000,
|
||
"chat": {"id": USER, "type": "private"},
|
||
"from": {"id": USER, "is_bot": False, "first_name": "h"},
|
||
"text": "see https://x.y",
|
||
"link_preview_options": {"url": "https://x.y"},
|
||
}
|
||
}
|
||
)
|
||
await stack.until(lambda: stack.sent_with("see https://x.y"), what="reply")
|
||
async with stack.world.db.session() as session:
|
||
rows = list((await session.exec(select(TelegramUpdate))).all())
|
||
preview = rows[0].payload["message"]["link_preview_options"]
|
||
assert preview["url"] == "https://x.y"
|
||
assert preview.get("is_disabled") is None
|
||
|
||
|
||
def _photo(caption: str) -> dict[str, Any]:
|
||
return {
|
||
"caption": caption,
|
||
"photo": [
|
||
{
|
||
"file_id": "f1",
|
||
"file_unique_id": "u1",
|
||
"width": 1,
|
||
"height": 1,
|
||
"file_size": 2048,
|
||
}
|
||
],
|
||
}
|
||
|
||
|
||
async def test_attachments_land_in_day_folder_with_keep_note(stack: Stack) -> None:
|
||
root = Path(tempfile.mkdtemp(prefix="beaver-att-"))
|
||
stack.tg.attachments = Attachments(dir=root, keep_days=3, tz="Asia/Bangkok")
|
||
stack.bot.media(_photo("смотри"))
|
||
await stack.until(lambda: stack.sent_with("смотри"), what="reply")
|
||
prompt = next(
|
||
p for c in ScriptedClient.instances for p in c.prompts if "смотри" in p
|
||
)
|
||
(day,) = list(root.iterdir())
|
||
assert day.name == stack.tg.attachments.day()
|
||
(saved,) = list(day.iterdir())
|
||
assert saved.read_bytes() == b"data" and saved.name.endswith("-u1.jpg")
|
||
assert prompt.endswith(
|
||
f"смотри\n\n[attachment: photo {saved}, 2 KB; kept 3 days, "
|
||
"move it into the vault if it matters]"
|
||
)
|
||
stack.tg.attachments = Attachments(dir=root, keep_days=None)
|
||
stack.bot.media(_photo("ещё"))
|
||
await stack.until(lambda: stack.sent_with("ещё"), what="second reply")
|
||
prompt = next(p for c in ScriptedClient.instances for p in c.prompts if "ещё" in p)
|
||
assert prompt.endswith(" KB]") and "kept" not in prompt
|
||
|
||
|
||
async def test_sticker_alone_becomes_an_unsupported_note(stack: Stack) -> None:
|
||
stack.bot.media(
|
||
{
|
||
"sticker": {
|
||
"file_id": "s1",
|
||
"file_unique_id": "su1",
|
||
"type": "regular",
|
||
"width": 1,
|
||
"height": 1,
|
||
"is_animated": False,
|
||
"is_video": False,
|
||
}
|
||
}
|
||
)
|
||
await stack.until(lambda: stack.sent_with("sticker"), what="reply")
|
||
prompt = next(
|
||
p for c in ScriptedClient.instances for p in c.prompts if "sticker" in p
|
||
)
|
||
assert prompt.endswith("[attachment: sticker - not supported]")
|
||
|
||
|
||
def test_sweep_drops_old_files_and_empty_day_folders() -> None:
|
||
root = Path(tempfile.mkdtemp(prefix="beaver-sweep-"))
|
||
old = root / "2020-01-01" / "1-old.jpg"
|
||
old.parent.mkdir()
|
||
old.write_bytes(b"x")
|
||
import os
|
||
|
||
os.utime(old, (1_577_836_800, 1_577_836_800))
|
||
fresh = root / "2099-01-01" / "2-new.jpg"
|
||
fresh.parent.mkdir()
|
||
fresh.write_bytes(b"y")
|
||
tg = TelegramFrontend(
|
||
token="t", user_id=USER, attachments=Attachments(dir=root, keep_days=None)
|
||
)
|
||
tg._sweep_attachments() # noqa: SLF001
|
||
assert old.exists()
|
||
tg.attachments = Attachments(dir=root, keep_days=7)
|
||
tg._sweep_attachments() # noqa: SLF001
|
||
assert not old.parent.exists() and fresh.exists()
|
||
|
||
|
||
async def test_topic_rename_updates_the_title(stack: Stack) -> None:
|
||
stack.bot.topic_created("план", 7)
|
||
stack.bot.message("hello", thread=7)
|
||
await stack.until(lambda: stack.sent_with("hello"), what="reply")
|
||
stack.bot.topic_edited("новое имя", 7)
|
||
|
||
async def titled() -> Any:
|
||
conv = await stack.world.conversations.find_bound(
|
||
frontend="telegram", external_id=f"{USER}/7"
|
||
)
|
||
return conv if conv is not None and conv.title == "новое имя" else None
|
||
|
||
deadline = asyncio.get_running_loop().time() + 5
|
||
branch = None
|
||
while branch is None and asyncio.get_running_loop().time() < deadline:
|
||
branch = await titled()
|
||
await asyncio.sleep(0.02)
|
||
assert branch is not None
|
||
stack.tg._topic_names.clear() # noqa: SLF001
|
||
assert await stack.tg.mark_topic(branch)
|
||
assert stack.bot.topics[-1] == "edit:7:✅ новое имя"
|
||
|
||
|
||
async def test_unknown_callback_strips_the_keyboard(stack: Stack) -> None:
|
||
stack.bot.callback("q:nope:0:1", 555)
|
||
await stack.until(
|
||
lambda: {"message_id": 555, "markup": None} in stack.bot.edits,
|
||
what="keyboard removed",
|
||
)
|
||
|
||
|
||
async def test_failed_background_turn_is_reported_once(stack: Stack) -> None:
|
||
stack.bot.message("hi")
|
||
await stack.until(lambda: stack.sent_with("hi"), what="reply")
|
||
master = await stack.world.conversations.find_bound(
|
||
frontend="telegram", external_id=GENERAL
|
||
)
|
||
event = {
|
||
"type": "turn.end",
|
||
"conversation_id": master.external_id,
|
||
"turn_id": "t-err",
|
||
"origin": "крон",
|
||
"stop": "error",
|
||
}
|
||
await stack.tg._on_event(event) # noqa: SLF001
|
||
await stack.tg._on_event(event) # noqa: SLF001
|
||
await stack.until(
|
||
lambda: stack.sent_with("a background turn (крон) failed"), what="notice"
|
||
)
|
||
await asyncio.sleep(0.2)
|
||
assert len([m for m in stack.bot.sent if "failed" in m["text"]]) == 1
|
||
await stack.tg._on_event({**event, "turn_id": "t-user", "origin": "user"}) # noqa: SLF001
|
||
notice = await stack.until(
|
||
lambda: stack.sent_with("⚠️ the turn failed"), what="user"
|
||
)
|
||
assert "background" not in notice["text"]
|
||
|
||
|
||
async def test_gone_topic_unbinds_and_the_next_message_starts_afresh(
|
||
stack: Stack,
|
||
) -> None:
|
||
stack.bot.message("hello", thread=7)
|
||
await stack.until(lambda: stack.sent_with("hello"), what="reply")
|
||
branch = await stack.world.conversations.find_bound(
|
||
frontend="telegram", external_id=f"{USER}/7"
|
||
)
|
||
stack.bot.gone_threads.add(7)
|
||
await stack.world.conversations.say(branch, "psst")
|
||
|
||
async def unbound() -> bool:
|
||
rows = await stack.world.conversations.bindings(branch)
|
||
return all(not b.visible for b in rows)
|
||
|
||
deadline = asyncio.get_running_loop().time() + 5
|
||
while not await unbound() and asyncio.get_running_loop().time() < deadline:
|
||
await asyncio.sleep(0.02)
|
||
assert await unbound()
|
||
assert branch.external_id not in stack.tg._targets # noqa: SLF001
|
||
async with stack.world.db.session() as session:
|
||
rows = list((await session.exec(select(Delivery))).all())
|
||
assert [r.status for r in rows if r.text == "psst"] == ["failed"]
|
||
stack.bot.gone_threads.clear()
|
||
stack.bot.message("again", thread=7)
|
||
await stack.until(lambda: stack.sent_with("again"), what="fresh reply")
|
||
fresh = await stack.world.conversations.find_bound(
|
||
frontend="telegram", external_id=f"{USER}/7"
|
||
)
|
||
assert fresh is not None and fresh.id != branch.id and fresh.status == "open"
|
||
|
||
|
||
async def test_photo_reaches_the_model_as_an_image_block(stack: Stack) -> None:
|
||
stack.bot.media(_photo("что на фото"))
|
||
await stack.until(lambda: stack.sent_with("что на фото"), what="reply")
|
||
client = next(c for c in ScriptedClient.instances if c.contents)
|
||
content = client.contents[-1]
|
||
assert content[0]["type"] == "text" and "что на фото" in content[0]["text"]
|
||
image = content[1]
|
||
assert image["type"] == "image"
|
||
assert image["source"]["media_type"] == "image/jpeg"
|
||
assert base64.b64decode(image["source"]["data"]) == b"data"
|
||
|
||
|
||
async def test_album_becomes_one_turn(stack: Stack) -> None:
|
||
stack.tg.album_delay = 0.2
|
||
stack.bot.media({**_photo("альбом"), "media_group_id": "g1"})
|
||
second = _photo("")
|
||
second["photo"][0].update(file_id="f2", file_unique_id="u2")
|
||
del second["caption"]
|
||
stack.bot.media({**second, "media_group_id": "g1"})
|
||
await stack.until(lambda: stack.sent_with("альбом"), what="reply")
|
||
await asyncio.sleep(0.3)
|
||
prompts = [
|
||
p for c in ScriptedClient.instances for p in c.prompts if "attachment" in p
|
||
]
|
||
assert len(prompts) == 1
|
||
assert prompts[0].count("[attachment: photo") == 2
|
||
client = next(c for c in ScriptedClient.instances if c.contents)
|
||
assert [b["type"] for b in client.contents[-1]] == ["text", "image", "image"]
|
||
|
||
|
||
async def test_stop_button_interrupts_the_turn(stack: Stack) -> None:
|
||
stack.bot.message("hi")
|
||
await stack.until(lambda: stack.sent_with("hi"), what="reply")
|
||
ScriptedClient.hold = asyncio.Event()
|
||
stack.bot.message("думай долго")
|
||
await stack.until(lambda: stack.tg._drafts, what="draft") # noqa: SLF001
|
||
draft = next(iter(stack.tg._drafts.values())) # noqa: SLF001
|
||
await stack.until(
|
||
lambda: any(d.get("draft_id") == draft.draft_id for d in stack.bot.drafts),
|
||
what="draft pushed",
|
||
)
|
||
assert stack.bot.drafts[-1]["can_stop"] is True
|
||
stack.bot.push(
|
||
{
|
||
"stopped_message_generation": {
|
||
"chat": {"id": USER, "type": "private"},
|
||
"draft_id": draft.draft_id,
|
||
"message_thread_id": stack.bot.drafts[-1].get("message_thread_id"),
|
||
}
|
||
}
|
||
)
|
||
await stack.until(lambda: stack.sent_with("думай долго"), what="reply after stop")
|
||
master = await stack.world.conversations.find_bound(
|
||
frontend="telegram", external_id=GENERAL
|
||
)
|
||
items = await stack.world.conversations.queue.recent(master.id, limit=1)
|
||
assert items[0].status == "interrupted"
|
||
ScriptedClient.hold = None
|
||
|
||
|
||
async def test_question_survives_a_gateway_restart(stack: Stack) -> None:
|
||
stack.bot.message("hi")
|
||
await stack.until(lambda: stack.sent_with("hi"), what="reply")
|
||
master = await stack.world.conversations.find_bound(
|
||
frontend="telegram", external_id=GENERAL
|
||
)
|
||
payload = {
|
||
"questions": [
|
||
{
|
||
"header": "Цвет",
|
||
"question": "Какой цвет?",
|
||
"options": [{"label": "Синий"}, {"label": "Красный"}],
|
||
"multiSelect": False,
|
||
}
|
||
]
|
||
}
|
||
asking = asyncio.create_task(
|
||
stack.world.conversations.ask(master.external_id, payload)
|
||
)
|
||
question = await stack.until(
|
||
lambda: stack.sent_with("Какой цвет?"), what="question"
|
||
)
|
||
pending = stack.world.conversations.pending_question(master.external_id)
|
||
assert pending is not None
|
||
row = await stack.world.conversations.get(master.external_id)
|
||
assert row.flags["ask"]["id"] == pending[0]
|
||
asking.cancel()
|
||
with contextlib.suppress(asyncio.CancelledError):
|
||
await asking
|
||
stack.tg._asks.clear() # noqa: SLF001
|
||
stack.world.conversations._questions.clear() # noqa: SLF001
|
||
stack.bot.callback(f"q:{pending[0]}:0:1", question["message_id"])
|
||
await stack.until(lambda: stack.sent_with("ok:Красный"), what="answer as message")
|
||
row = await stack.world.conversations.get(master.external_id)
|
||
assert row.flags.get("ask") is None
|
||
|
||
|
||
def _sendable(name: str, body: bytes = b"x") -> tuple[Path, Path]:
|
||
root = Path(tempfile.mkdtemp(prefix="beaver-send-")).resolve()
|
||
file = root / name
|
||
file.write_bytes(body)
|
||
return root, file
|
||
|
||
|
||
def test_send_files_takes_only_paths_inside_its_roots() -> None:
|
||
root, note = _sendable("note.txt")
|
||
outside = Path(tempfile.mkdtemp(prefix="beaver-out-")).resolve()
|
||
(outside / "secret.txt").write_bytes(b"y")
|
||
(root / "escape.txt").symlink_to(outside / "secret.txt")
|
||
files = SendFiles(roots=(root,))
|
||
|
||
assert files.resolve(f"{root}/../{root.name}/./note.txt") == note
|
||
|
||
for bad in (
|
||
str(outside / "secret.txt"),
|
||
f"{root}/../{outside.name}/secret.txt",
|
||
str(root / "escape.txt"),
|
||
):
|
||
with pytest.raises(SendFileError, match="outside the allowed roots"):
|
||
files.resolve(bad)
|
||
with pytest.raises(SendFileError, match="is not a file"):
|
||
files.resolve(str(root / "gone.txt"))
|
||
with pytest.raises(SendFileError, match="outside the allowed roots"):
|
||
SendFiles().resolve(str(note))
|
||
|
||
|
||
def test_send_files_stops_a_file_over_the_limit() -> None:
|
||
root, big = _sendable("big.bin", b"0" * (2 * 1024 * 1024))
|
||
files = SendFiles(roots=(root,), max_bytes=1024 * 1024)
|
||
with pytest.raises(SendFileError, match="big.bin is 2 MB, over the 1 MB limit"):
|
||
files.resolve(str(big))
|
||
assert SendFiles(roots=(root,), max_bytes=4 * 1024 * 1024).resolve(str(big)) == big
|
||
assert files.note == f"paths under {root}, up to 1 MB"
|
||
|
||
|
||
def test_method_follows_the_extension_unless_it_is_given() -> None:
|
||
picked = {
|
||
name: method_for(Path(name))
|
||
for name in ("a.JPG", "b.gif", "c.mp3", "d.ogg", "e.mov", "f.zip", "g.svg")
|
||
}
|
||
assert picked == {
|
||
"a.JPG": "photo",
|
||
"b.gif": "animation",
|
||
"c.mp3": "audio",
|
||
"d.ogg": "voice",
|
||
"e.mov": "video",
|
||
"f.zip": "document",
|
||
"g.svg": "photo",
|
||
}
|
||
assert method_for(Path("a.jpg"), "document") == "document"
|
||
with pytest.raises(SendFileError, match="unknown method 'sticker'"):
|
||
method_for(Path("a.jpg"), "sticker")
|
||
|
||
|
||
async def _send(stack: Stack, conv, **args: Any) -> dict[str, Any]:
|
||
tool = next(
|
||
t
|
||
for t in _tools(stack.world.conversations, conv.external_id)
|
||
if t.name == "send_file"
|
||
)
|
||
return await tool.handler(args)
|
||
|
||
|
||
async def test_send_file_lands_in_the_topic_of_the_conversation(stack: Stack) -> None:
|
||
stack.bot.message("work", thread=11)
|
||
await stack.until(lambda: stack.sent_with("work"), what="branch reply")
|
||
branch = await stack.world.conversations.find_bound(
|
||
frontend="telegram", external_id=f"{USER}/11"
|
||
)
|
||
root, shot = _sendable("график.png", b"png")
|
||
stack.tg.send_files = SendFiles(roots=(root,))
|
||
|
||
result = await _send(stack, branch, path=str(shot), caption="**вот**")
|
||
|
||
assert result["content"][0]["text"].startswith("sent график.png as photo")
|
||
(sent,) = stack.bot.files
|
||
assert sent["method"] == "photo"
|
||
assert sent["message_thread_id"] == 11
|
||
assert sent["photo"].path == shot
|
||
assert sent["caption"] == "<b>вот</b>" and sent["parse_mode"] == "HTML"
|
||
|
||
|
||
async def test_long_caption_follows_the_file_as_a_message(stack: Stack) -> None:
|
||
stack.bot.message("hi")
|
||
await stack.until(lambda: stack.sent_with("hi"), what="reply")
|
||
master = await stack.world.conversations.find_bound(
|
||
frontend="telegram", external_id=GENERAL
|
||
)
|
||
root, doc = _sendable("отчёт.pdf")
|
||
stack.tg.send_files = SendFiles(roots=(root,))
|
||
|
||
await _send(stack, master, path=str(doc), caption="долго " * 300)
|
||
|
||
(sent,) = stack.bot.files
|
||
assert sent["method"] == "document" and "caption" not in sent
|
||
await stack.until(lambda: stack.sent_with("долго"), what="caption message")
|
||
|
||
|
||
async def test_send_file_without_a_window_falls_back_to_the_master(
|
||
stack: Stack,
|
||
) -> None:
|
||
stack.bot.message("hi")
|
||
await stack.until(lambda: stack.sent_with("hi"), what="reply")
|
||
job = await stack.world.conversations.spawn(
|
||
kind="job", agent="a", seed="brief", text="job", parent=None, origin="test"
|
||
)
|
||
root, song = _sendable("Пикник.mp3")
|
||
stack.tg.send_files = SendFiles(roots=(root,))
|
||
|
||
await _send(stack, job, path=str(song))
|
||
|
||
(sent,) = stack.bot.files
|
||
assert sent["method"] == "audio" and sent["title"] == "Пикник"
|
||
assert (sent["chat_id"], sent["message_thread_id"]) == (USER, 901)
|
||
|
||
|
||
async def test_a_refused_upload_comes_back_to_the_agent(stack: Stack) -> None:
|
||
stack.bot.message("hi")
|
||
await stack.until(lambda: stack.sent_with("hi"), what="reply")
|
||
master = await stack.world.conversations.find_bound(
|
||
frontend="telegram", external_id=GENERAL
|
||
)
|
||
root, shot = _sendable("bad.png")
|
||
stack.tg.send_files = SendFiles(roots=(root,))
|
||
stack.bot.reject_files = True
|
||
|
||
result = await _send(stack, master, path=str(shot))
|
||
|
||
assert result["is_error"]
|
||
assert "Telegram refused bad.png as photo" in result["content"][0]["text"]
|
||
|
||
outside = await _send(stack, master, path="/etc/passwd")
|
||
assert outside["is_error"]
|
||
assert "outside the allowed roots" in outside["content"][0]["text"]
|