main
py 96 lines 2.59 KB
Raw
1 import asyncio
2
3 from plugins._telegram_integration.extensions.python._functions.agent.Agent.handle_exception.end import (
4 _85_telegram_error,
5 )
6 from plugins._telegram_integration.helpers import error_ui
7 from plugins._telegram_integration.helpers.constants import (
8 CTX_TG_BOT,
9 CTX_TG_ERROR_SENT,
10 CTX_TG_REPLY_TO,
11 CTX_TG_TYPING_STOP,
12 )
13
14
15 class FakeContext:
16 def __init__(self):
17 self.data = {
18 CTX_TG_BOT: "main",
19 CTX_TG_REPLY_TO: 456,
20 }
21
22
23 class FakeAgent:
24 number = 0
25
26 def __init__(self):
27 self.context = FakeContext()
28
29
30 def test_friendly_error_message_for_missing_api_key():
31 message = error_ui.friendly_error_message(
32 RuntimeError("OPENAI_API_KEY is missing or invalid")
33 )
34
35 assert "provider setup issue" in message
36 assert "API key" in message
37 assert "OPENAI_API_KEY is missing or invalid" in message
38
39
40 def test_friendly_error_message_for_rate_limit():
41 message = error_ui.friendly_error_message(
42 RuntimeError("Rate limit exceeded: too many requests")
43 )
44
45 assert "rate limited" in message
46 assert "too many requests" in message
47
48
49 def test_telegram_exception_hook_sends_once_and_cleans_stream_state(monkeypatch):
50 sends = []
51 cleared = []
52 typing_stopped = []
53
54 async def fake_send(context, text, attachments=None, keyboard=None):
55 sends.append(
56 {
57 "text": text,
58 "reply_to": context.data.get(CTX_TG_REPLY_TO),
59 "attachments": attachments,
60 "keyboard": keyboard,
61 }
62 )
63 return None
64
65 def fake_clear(context):
66 cleared.append(True)
67
68 class FakeStop:
69 def set(self):
70 typing_stopped.append(True)
71
72 agent = FakeAgent()
73 agent.context.data[CTX_TG_TYPING_STOP] = FakeStop()
74
75 monkeypatch.setattr(
76 "plugins._telegram_integration.helpers.handler.send_telegram_reply",
77 fake_send,
78 )
79 monkeypatch.setattr(
80 "plugins._telegram_integration.helpers.draft_stream.clear",
81 fake_clear,
82 )
83
84 extension = _85_telegram_error.TelegramFriendlyError(agent=agent)
85 data = {"exception": RuntimeError("provider returned 503 service unavailable")}
86
87 asyncio.run(extension.execute(data=data))
88 asyncio.run(extension.execute(data=data))
89
90 assert len(sends) == 1
91 assert "could not complete the model request" in sends[0]["text"]
92 assert sends[0]["reply_to"] == 456
93 assert cleared == [True]
94 assert typing_stopped == [True]
95 assert agent.context.data[CTX_TG_ERROR_SENT] is True
96 assert CTX_TG_REPLY_TO not in agent.context.data