@samitouri / QOSami_Agents_Zero-1
Code Commits Issues Pulls Sessions CI/CD
main
main
py 253 lines 7.22 KB
Raw
1 from __future__ import annotations
2
3 import asyncio
4 from types import SimpleNamespace
5 from unittest.mock import AsyncMock
6
7 import pytest
8
9 from helpers.ws_manager import WsResult
10 from plugins._a0_connector.api import ws_connector as ws_module
11 from plugins._a0_connector.api.ws_connector import (
12 WsConnector,
13 _attachment_log_metadata,
14 )
15 from plugins._a0_connector.helpers.event_bridge import log_entry_to_connector_event
16
17
18 def test_attachment_log_metadata_keeps_only_safe_basenames() -> None:
19 assert _attachment_log_metadata(
20 [
21 "/a0/usr/uploads/scan.png",
22 r"C:\\Users\\person\\result.jpg",
23 "https://agent.test/api/image_get?path=/a0/usr/uploads/chart.webp&token=secret#view",
24 "/a0/usr/uploads/",
25 "",
26 ]
27 ) == {
28 "attachments": ["scan.png", "result.jpg", "image_get"]
29 }
30
31
32 def test_attachment_log_metadata_omits_empty_metadata() -> None:
33 assert _attachment_log_metadata([]) == {}
34 assert _attachment_log_metadata(["", "/"]) == {}
35
36
37 def test_attachment_log_metadata_decodes_encoded_separators() -> None:
38 assert _attachment_log_metadata(
39 [
40 "https://host/%2Fhome%2Falice%2Fsecret.png",
41 "https://host/C:%5CUsers%5CAlice%5Csecret.png",
42 ]
43 ) == {"attachments": ["secret.png", "secret.png"]}
44
45
46 def test_attachment_log_metadata_strips_encoded_query_and_fragment_suffixes() -> None:
47 assert _attachment_log_metadata(
48 [
49 "https://host/report%3Ftoken%3Dsecret.png",
50 "https://host/image%23private-fragment.png",
51 ]
52 ) == {"attachments": ["report", "image"]}
53
54
55 def test_attachment_log_metadata_decodes_double_encoded_separators() -> None:
56 assert _attachment_log_metadata(
57 [
58 "https://host/%252Fhome%252Fuser%252Fsecret.png",
59 "https://host/C:%255CUsers%255CAlice%255Csecret.png",
60 ]
61 ) == {"attachments": ["secret.png", "secret.png"]}
62
63
64 def test_attachment_log_metadata_strips_double_encoded_delimiter_suffixes() -> None:
65 assert _attachment_log_metadata(
66 [
67 "https://host/report%253Ftoken%253Dredacted.png",
68 "https://host/image%2523private-fragment.png",
69 ]
70 ) == {"attachments": ["report", "image"]}
71
72
73 def test_attachment_log_metadata_omits_paths_exceeding_decode_limit() -> None:
74 assert _attachment_log_metadata(
75 ["https://host/%25252Fhome%25252Fuser%25252Fsecret.png"]
76 ) == {}
77
78
79 class RecordingLog:
80 def __init__(self) -> None:
81 self.calls: list[dict[str, object]] = []
82
83 def log(self, **kwargs: object) -> None:
84 self.calls.append(dict(kwargs))
85
86
87 @pytest.mark.asyncio
88 async def test_websocket_attachment_names_reach_replayed_user_event(
89 monkeypatch: pytest.MonkeyPatch,
90 ) -> None:
91 log = RecordingLog()
92 context = SimpleNamespace(log=log)
93 handler = WsConnector(None, None)
94 monkeypatch.setattr(
95 handler,
96 "_resolve_context",
97 AsyncMock(return_value=(context, "ctx-1")),
98 )
99 monkeypatch.setattr(
100 ws_module,
101 "subscribed_contexts_for_sid",
102 lambda sid: {"ctx-1"} if sid == "sid-cli" else set(),
103 )
104
105 scheduled: list[bool] = []
106
107 def close_scheduled(coroutine: object) -> SimpleNamespace:
108 close = getattr(coroutine, "close")
109 close()
110 scheduled.append(True)
111 return SimpleNamespace()
112
113 monkeypatch.setattr(asyncio, "create_task", close_scheduled)
114
115 result = await handler._handle_send_message(
116 {
117 "context_id": "ctx-1",
118 "message": "Review these",
119 "attachments": [
120 "/a0/usr/uploads/scan.png",
121 "/a0/usr/uploads/result.jpg",
122 ],
123 "client_message_id": "client-1",
124 },
125 "sid-cli",
126 )
127
128 assert result == {
129 "context_id": "ctx-1",
130 "status": "accepted",
131 "client_message_id": "client-1",
132 }
133 assert scheduled == [True]
134 assert log.calls == [
135 {
136 "type": "user",
137 "heading": "",
138 "content": "Review these",
139 "kvps": {"attachments": ["scan.png", "result.jpg"]},
140 "id": "client-1",
141 }
142 ]
143
144 replayed = log_entry_to_connector_event(
145 {"no": 0, **log.calls[0]},
146 "ctx-1",
147 )
148 assert replayed["event"] == "user_message"
149 assert replayed["data"]["meta"] == {
150 "attachments": ["scan.png", "result.jpg"]
151 }
152
153
154 @pytest.mark.asyncio
155 async def test_websocket_text_only_message_keeps_empty_kvps(
156 monkeypatch: pytest.MonkeyPatch,
157 ) -> None:
158 log = RecordingLog()
159 context = SimpleNamespace(log=log)
160 handler = WsConnector(None, None)
161 monkeypatch.setattr(
162 handler,
163 "_resolve_context",
164 AsyncMock(return_value=(context, "ctx-1")),
165 )
166 monkeypatch.setattr(
167 ws_module,
168 "subscribed_contexts_for_sid",
169 lambda sid: {"ctx-1"} if sid == "sid-cli" else set(),
170 )
171
172 def close_scheduled(coroutine: object) -> SimpleNamespace:
173 getattr(coroutine, "close")()
174 return SimpleNamespace()
175
176 monkeypatch.setattr(asyncio, "create_task", close_scheduled)
177 await handler._handle_send_message(
178 {"context_id": "ctx-1", "message": "Text only"},
179 "sid-cli",
180 )
181 assert log.calls[0]["kvps"] == {}
182
183
184 @pytest.mark.asyncio
185 async def test_websocket_malformed_attachment_url_keeps_message_delivery(
186 monkeypatch: pytest.MonkeyPatch,
187 ) -> None:
188 log = RecordingLog()
189 context = SimpleNamespace(log=log)
190 handler = WsConnector(None, None)
191 monkeypatch.setattr(
192 handler,
193 "_resolve_context",
194 AsyncMock(return_value=(context, "ctx-1")),
195 )
196 monkeypatch.setattr(
197 ws_module,
198 "subscribed_contexts_for_sid",
199 lambda sid: {"ctx-1"} if sid == "sid-cli" else set(),
200 )
201
202 scheduled: list[bool] = []
203
204 def close_scheduled(coroutine: object) -> SimpleNamespace:
205 getattr(coroutine, "close")()
206 scheduled.append(True)
207 return SimpleNamespace()
208
209 monkeypatch.setattr(asyncio, "create_task", close_scheduled)
210
211 result = await handler._handle_send_message(
212 {
213 "context_id": "ctx-1",
214 "message": "Review this",
215 "attachments": ["http://["],
216 "client_message_id": "client-malformed",
217 },
218 "sid-cli",
219 )
220
221 assert result == {
222 "context_id": "ctx-1",
223 "status": "accepted",
224 "client_message_id": "client-malformed",
225 }
226 assert scheduled == [True]
227 assert log.calls[0]["kvps"] == {}
228
229
230 @pytest.mark.asyncio
231 @pytest.mark.parametrize(
232 ("payload", "code"),
233 [
234 (
235 {
236 "message": "image",
237 "attachments": [{"path": "data:image/png;base64,AAAA"}],
238 },
239 "INVALID_ATTACHMENTS",
240 ),
241 ({"message": "", "attachments": []}, "MISSING_MESSAGE"),
242 ],
243 )
244 async def test_websocket_rejected_attachments_do_not_reach_context(
245 payload: dict[str, object],
246 code: str,
247 ) -> None:
248 handler = WsConnector(None, None)
249 result = await handler.process("connector_send_message", payload, "sid-cli")
250 assert isinstance(result, WsResult)
251 rendered = result.as_result(handler_id="test", fallback_correlation_id=None)
252 assert rendered["ok"] is False
253 assert rendered["error"]["code"] == code