106 lines
2.7 KiB
Python
106 lines
2.7 KiB
Python
|
# retoor <retoor@molodetz.nl>
|
||
|
|
|
||
|
|
from devplacepy.services.telegram.bridge import TelegramConnection
|
||
|
|
from tests.conftest import run_async
|
||
|
|
|
||
|
|
|
||
|
|
class _Service:
|
||
|
|
def __init__(self):
|
||
|
|
self.sent = []
|
||
|
|
self.edits = []
|
||
|
|
|
||
|
|
async def send(self, chat_id, text, **kwargs):
|
||
|
|
self.sent.append(text)
|
||
|
|
return 100 + len(self.sent)
|
||
|
|
|
||
|
|
async def edit(self, chat_id, message_id, text, **kwargs):
|
||
|
|
self.edits.append(text)
|
||
|
|
return True
|
||
|
|
|
||
|
|
async def chat_action(self, chat_id, action="typing"):
|
||
|
|
return None
|
||
|
|
|
||
|
|
|
||
|
|
class _Session:
|
||
|
|
def resolve_query(self, query_id, payload):
|
||
|
|
pass
|
||
|
|
|
||
|
|
def resolve_interaction(self, interaction_id, payload):
|
||
|
|
pass
|
||
|
|
|
||
|
|
|
||
|
|
def _flow(payloads):
|
||
|
|
async def run():
|
||
|
|
import asyncio
|
||
|
|
|
||
|
|
service = _Service()
|
||
|
|
conn = TelegramConnection(service, _Session(), 7, 42)
|
||
|
|
try:
|
||
|
|
for payload in payloads:
|
||
|
|
await conn.send_json(payload)
|
||
|
|
finally:
|
||
|
|
conn.done.set()
|
||
|
|
conn._stop_typing()
|
||
|
|
await asyncio.gather(conn._typing_task, return_exceptions=True)
|
||
|
|
return service
|
||
|
|
|
||
|
|
return run_async(run())
|
||
|
|
|
||
|
|
|
||
|
|
def test_trace_without_agent_unchanged():
|
||
|
|
service = _flow([{"type": "trace", "event": "call", "name": "fetch_url"}])
|
||
|
|
assert service.edits == ["<i>Working... (fetch_url)</i>"]
|
||
|
|
assert service.sent == []
|
||
|
|
|
||
|
|
|
||
|
|
def test_trace_with_agent_labels_progress():
|
||
|
|
service = _flow(
|
||
|
|
[
|
||
|
|
{
|
||
|
|
"type": "trace",
|
||
|
|
"event": "call",
|
||
|
|
"name": "fetch_url",
|
||
|
|
"agent": "sg",
|
||
|
|
"uid": "uid-1",
|
||
|
|
}
|
||
|
|
]
|
||
|
|
)
|
||
|
|
assert service.edits == ["<i>Working... ([sg] fetch_url)</i>"]
|
||
|
|
assert service.sent == []
|
||
|
|
|
||
|
|
|
||
|
|
def test_agent_frame_while_open_edits_placeholder():
|
||
|
|
service = _flow(
|
||
|
|
[
|
||
|
|
{
|
||
|
|
"type": "agent",
|
||
|
|
"kind": "needs_user",
|
||
|
|
"uid": "uid-1",
|
||
|
|
"label": "sg",
|
||
|
|
"payload": "Budget?",
|
||
|
|
}
|
||
|
|
]
|
||
|
|
)
|
||
|
|
assert service.edits == ["<i>[sg] needs_user</i>"]
|
||
|
|
assert service.sent == []
|
||
|
|
|
||
|
|
|
||
|
|
def test_agent_frame_after_finalize_sends_new_message():
|
||
|
|
service = _flow(
|
||
|
|
[
|
||
|
|
{"type": "reply", "text": "On it, working in the background."},
|
||
|
|
{
|
||
|
|
"type": "agent",
|
||
|
|
"kind": "done",
|
||
|
|
"uid": "uid-1",
|
||
|
|
"label": "sg",
|
||
|
|
"payload": "The comparison is ready.",
|
||
|
|
},
|
||
|
|
]
|
||
|
|
)
|
||
|
|
assert len(service.edits) == 1
|
||
|
|
assert len(service.sent) == 1
|
||
|
|
assert "Subagent" in service.sent[0]
|
||
|
|
assert "sg" in service.sent[0]
|
||
|
|
assert "finished" in service.sent[0]
|