| 1 | import asyncio |
| 2 | |
| 3 | from plugins._telegram_integration.helpers import heartbeat |
| 4 | from plugins._telegram_integration.helpers.constants import ( |
| 5 | CTX_TG_BOT, |
| 6 | CTX_TG_CHAT_ID, |
| 7 | CTX_TG_HEARTBEAT_STOP, |
| 8 | CTX_TG_HEARTBEAT_TASK, |
| 9 | ) |
| 10 | |
| 11 | |
| 12 | class FakeBot: |
| 13 | class Bot: |
| 14 | token = "token" |
| 15 | |
| 16 | bot = Bot() |
| 17 | |
| 18 | |
| 19 | class FakeLog: |
| 20 | progress = "icon://search[Search] A0: Searching official docs" |
| 21 | |
| 22 | |
| 23 | class FakeContext: |
| 24 | def __init__(self): |
| 25 | self.data = { |
| 26 | CTX_TG_BOT: "main", |
| 27 | CTX_TG_CHAT_ID: 123, |
| 28 | } |
| 29 | self.log = FakeLog() |
| 30 | |
| 31 | |
| 32 | def test_heartbeat_text_omits_iteration_count(): |
| 33 | text = heartbeat.heartbeat_text(180, "waiting for provider response") |
| 34 | |
| 35 | assert text == "Still working... (3 min elapsed - waiting for provider response)" |
| 36 | assert "iteration" not in text |
| 37 | |
| 38 | |
| 39 | def test_heartbeat_sends_periodic_status_and_stops(monkeypatch): |
| 40 | calls = [] |
| 41 | |
| 42 | async def fake_send(token, chat_id, text, parse_mode=None, **kwargs): |
| 43 | calls.append( |
| 44 | { |
| 45 | "token": token, |
| 46 | "chat_id": chat_id, |
| 47 | "text": text, |
| 48 | "parse_mode": parse_mode, |
| 49 | } |
| 50 | ) |
| 51 | return len(calls) |
| 52 | |
| 53 | context = FakeContext() |
| 54 | monkeypatch.setattr(heartbeat, "HEARTBEAT_INTERVAL_SECONDS", 0.01) |
| 55 | monkeypatch.setattr(heartbeat, "get_bot", lambda name: FakeBot()) |
| 56 | monkeypatch.setattr(heartbeat.tc, "raw_send_text", fake_send) |
| 57 | |
| 58 | async def run(): |
| 59 | await heartbeat.start(context) |
| 60 | await asyncio.sleep(0.025) |
| 61 | await heartbeat.stop(context) |
| 62 | |
| 63 | asyncio.run(run()) |
| 64 | |
| 65 | assert calls |
| 66 | assert calls[0]["chat_id"] == 123 |
| 67 | assert calls[0]["parse_mode"] is None |
| 68 | assert "Still working..." in calls[0]["text"] |
| 69 | assert "A0: Searching official docs" in calls[0]["text"] |
| 70 | assert CTX_TG_HEARTBEAT_TASK not in context.data |
| 71 | assert CTX_TG_HEARTBEAT_STOP not in context.data |