proximl 0.5.17__py3-none-any.whl → 1.0.0__py3-none-any.whl
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- examples/local_storage.py +0 -2
- proximl/__init__.py +1 -1
- proximl/checkpoints.py +56 -57
- proximl/cli/__init__.py +6 -3
- proximl/cli/checkpoint.py +18 -57
- proximl/cli/dataset.py +17 -57
- proximl/cli/job/__init__.py +11 -53
- proximl/cli/job/create.py +51 -24
- proximl/cli/model.py +14 -56
- proximl/cli/volume.py +18 -57
- proximl/datasets.py +50 -55
- proximl/jobs.py +239 -68
- proximl/models.py +51 -55
- proximl/proximl.py +50 -16
- proximl/utils/__init__.py +1 -0
- proximl/{auth.py → utils/auth.py} +4 -3
- proximl/utils/transfer.py +587 -0
- proximl/volumes.py +48 -53
- {proximl-0.5.17.dist-info → proximl-1.0.0.dist-info}/METADATA +3 -3
- {proximl-0.5.17.dist-info → proximl-1.0.0.dist-info}/RECORD +52 -50
- tests/integration/test_checkpoints_integration.py +4 -3
- tests/integration/test_datasets_integration.py +5 -3
- tests/integration/test_jobs_integration.py +33 -27
- tests/integration/test_models_integration.py +7 -3
- tests/integration/test_volumes_integration.py +2 -2
- tests/unit/cli/test_cli_checkpoint_unit.py +312 -1
- tests/unit/cloudbender/test_nodes_unit.py +112 -0
- tests/unit/cloudbender/test_providers_unit.py +96 -0
- tests/unit/cloudbender/test_regions_unit.py +106 -0
- tests/unit/cloudbender/test_services_unit.py +141 -0
- tests/unit/conftest.py +23 -10
- tests/unit/projects/test_project_data_connectors_unit.py +39 -0
- tests/unit/projects/test_project_datastores_unit.py +37 -0
- tests/unit/projects/test_project_members_unit.py +46 -0
- tests/unit/projects/test_project_services_unit.py +65 -0
- tests/unit/projects/test_projects_unit.py +16 -0
- tests/unit/test_auth_unit.py +17 -2
- tests/unit/test_checkpoints_unit.py +256 -71
- tests/unit/test_datasets_unit.py +218 -68
- tests/unit/test_exceptions.py +133 -0
- tests/unit/test_gpu_types_unit.py +11 -1
- tests/unit/test_jobs_unit.py +1014 -95
- tests/unit/test_main_unit.py +20 -0
- tests/unit/test_models_unit.py +218 -70
- tests/unit/test_proximl_unit.py +627 -3
- tests/unit/test_volumes_unit.py +211 -70
- tests/unit/utils/__init__.py +1 -0
- tests/unit/utils/test_transfer_unit.py +4260 -0
- proximl/cli/connection.py +0 -61
- proximl/connections.py +0 -621
- tests/unit/test_connections_unit.py +0 -182
- {proximl-0.5.17.dist-info → proximl-1.0.0.dist-info}/LICENSE +0 -0
- {proximl-0.5.17.dist-info → proximl-1.0.0.dist-info}/WHEEL +0 -0
- {proximl-0.5.17.dist-info → proximl-1.0.0.dist-info}/entry_points.txt +0 -0
- {proximl-0.5.17.dist-info → proximl-1.0.0.dist-info}/top_level.txt +0 -0
|
@@ -0,0 +1,4260 @@
|
|
|
1
|
+
import os
|
|
2
|
+
import re
|
|
3
|
+
import asyncio
|
|
4
|
+
import tempfile
|
|
5
|
+
from unittest.mock import (
|
|
6
|
+
Mock,
|
|
7
|
+
AsyncMock,
|
|
8
|
+
patch,
|
|
9
|
+
mock_open,
|
|
10
|
+
MagicMock,
|
|
11
|
+
)
|
|
12
|
+
from pytest import mark, fixture, raises
|
|
13
|
+
from aiohttp import ClientResponseError, ClientSession
|
|
14
|
+
from aiohttp.client_exceptions import (
|
|
15
|
+
ClientConnectorError,
|
|
16
|
+
ClientPayloadError,
|
|
17
|
+
ServerTimeoutError,
|
|
18
|
+
ServerDisconnectedError,
|
|
19
|
+
ClientOSError,
|
|
20
|
+
InvalidURL,
|
|
21
|
+
)
|
|
22
|
+
|
|
23
|
+
import proximl.utils.transfer as specimen
|
|
24
|
+
from proximl.exceptions import ConnectionError, ProxiMLException
|
|
25
|
+
|
|
26
|
+
pytestmark = [mark.sdk, mark.unit]
|
|
27
|
+
|
|
28
|
+
|
|
29
|
+
class NormalizeEndpointTests:
|
|
30
|
+
def test_normalize_endpoint_with_https(self):
|
|
31
|
+
result = specimen.normalize_endpoint("https://example.com")
|
|
32
|
+
assert result == "https://example.com"
|
|
33
|
+
|
|
34
|
+
def test_normalize_endpoint_with_http(self):
|
|
35
|
+
result = specimen.normalize_endpoint("http://example.com")
|
|
36
|
+
assert result == "http://example.com"
|
|
37
|
+
|
|
38
|
+
def test_normalize_endpoint_without_protocol(self):
|
|
39
|
+
result = specimen.normalize_endpoint("example.com")
|
|
40
|
+
assert result == "https://example.com"
|
|
41
|
+
|
|
42
|
+
def test_normalize_endpoint_with_trailing_slash(self):
|
|
43
|
+
result = specimen.normalize_endpoint("https://example.com/")
|
|
44
|
+
assert result == "https://example.com"
|
|
45
|
+
|
|
46
|
+
def test_normalize_endpoint_empty_string(self):
|
|
47
|
+
with raises(ValueError, match="Endpoint URL cannot be empty"):
|
|
48
|
+
specimen.normalize_endpoint("")
|
|
49
|
+
|
|
50
|
+
def test_normalize_endpoint_multiple_trailing_slashes(self):
|
|
51
|
+
result = specimen.normalize_endpoint("https://example.com///")
|
|
52
|
+
assert result == "https://example.com"
|
|
53
|
+
|
|
54
|
+
|
|
55
|
+
class RetryRequestTests:
|
|
56
|
+
@mark.asyncio
|
|
57
|
+
async def test_retry_request_success_first_attempt(self):
|
|
58
|
+
func = AsyncMock(return_value="success")
|
|
59
|
+
result = await specimen.retry_request(func)
|
|
60
|
+
assert result == "success"
|
|
61
|
+
assert func.call_count == 1
|
|
62
|
+
|
|
63
|
+
@mark.asyncio
|
|
64
|
+
async def test_retry_request_retry_on_502(self):
|
|
65
|
+
func = AsyncMock(
|
|
66
|
+
side_effect=[
|
|
67
|
+
ClientResponseError(
|
|
68
|
+
request_info=Mock(),
|
|
69
|
+
history=(),
|
|
70
|
+
status=502,
|
|
71
|
+
message="Bad Gateway",
|
|
72
|
+
),
|
|
73
|
+
"success",
|
|
74
|
+
]
|
|
75
|
+
)
|
|
76
|
+
with patch("asyncio.sleep", new_callable=AsyncMock):
|
|
77
|
+
result = await specimen.retry_request(func, max_retries=3)
|
|
78
|
+
assert result == "success"
|
|
79
|
+
assert func.call_count == 2
|
|
80
|
+
|
|
81
|
+
@mark.asyncio
|
|
82
|
+
async def test_retry_request_retry_on_503(self):
|
|
83
|
+
func = AsyncMock(
|
|
84
|
+
side_effect=[
|
|
85
|
+
ClientResponseError(
|
|
86
|
+
request_info=Mock(),
|
|
87
|
+
history=(),
|
|
88
|
+
status=503,
|
|
89
|
+
message="Service Unavailable",
|
|
90
|
+
),
|
|
91
|
+
"success",
|
|
92
|
+
]
|
|
93
|
+
)
|
|
94
|
+
with patch("asyncio.sleep", new_callable=AsyncMock):
|
|
95
|
+
result = await specimen.retry_request(func, max_retries=3)
|
|
96
|
+
assert result == "success"
|
|
97
|
+
|
|
98
|
+
@mark.asyncio
|
|
99
|
+
async def test_retry_request_retry_on_504(self):
|
|
100
|
+
func = AsyncMock(
|
|
101
|
+
side_effect=[
|
|
102
|
+
ClientResponseError(
|
|
103
|
+
request_info=Mock(),
|
|
104
|
+
history=(),
|
|
105
|
+
status=504,
|
|
106
|
+
message="Gateway Timeout",
|
|
107
|
+
),
|
|
108
|
+
"success",
|
|
109
|
+
]
|
|
110
|
+
)
|
|
111
|
+
with patch("asyncio.sleep", new_callable=AsyncMock):
|
|
112
|
+
result = await specimen.retry_request(func, max_retries=3)
|
|
113
|
+
assert result == "success"
|
|
114
|
+
|
|
115
|
+
@mark.asyncio
|
|
116
|
+
async def test_retry_request_max_retries_exceeded(self):
|
|
117
|
+
func = AsyncMock(
|
|
118
|
+
side_effect=ClientResponseError(
|
|
119
|
+
request_info=Mock(),
|
|
120
|
+
history=(),
|
|
121
|
+
status=502,
|
|
122
|
+
message="Bad Gateway",
|
|
123
|
+
)
|
|
124
|
+
)
|
|
125
|
+
with patch("asyncio.sleep", new_callable=AsyncMock):
|
|
126
|
+
with raises(ClientResponseError):
|
|
127
|
+
await specimen.retry_request(func, max_retries=3)
|
|
128
|
+
assert func.call_count == 3
|
|
129
|
+
|
|
130
|
+
@mark.asyncio
|
|
131
|
+
async def test_retry_request_non_retry_status(self):
|
|
132
|
+
func = AsyncMock(
|
|
133
|
+
side_effect=ClientResponseError(
|
|
134
|
+
request_info=Mock(),
|
|
135
|
+
history=(),
|
|
136
|
+
status=400,
|
|
137
|
+
message="Bad Request",
|
|
138
|
+
)
|
|
139
|
+
)
|
|
140
|
+
with patch("asyncio.sleep", new_callable=AsyncMock):
|
|
141
|
+
with raises(ClientResponseError):
|
|
142
|
+
await specimen.retry_request(func, max_retries=3)
|
|
143
|
+
assert func.call_count == 1
|
|
144
|
+
|
|
145
|
+
@mark.asyncio
|
|
146
|
+
async def test_retry_request_connection_error(self):
|
|
147
|
+
func = AsyncMock(
|
|
148
|
+
side_effect=ClientConnectorError(
|
|
149
|
+
connection_key=Mock(), os_error=OSError("Connection failed")
|
|
150
|
+
)
|
|
151
|
+
)
|
|
152
|
+
with patch("asyncio.sleep", new_callable=AsyncMock):
|
|
153
|
+
with raises(ClientConnectorError):
|
|
154
|
+
await specimen.retry_request(func, max_retries=2)
|
|
155
|
+
# DNS errors now use DNS_MAX_RETRIES (7) instead of max_retries when ClientConnectorError is encountered
|
|
156
|
+
assert func.call_count == 7
|
|
157
|
+
|
|
158
|
+
@mark.asyncio
|
|
159
|
+
async def test_retry_request_server_disconnected_error(self):
|
|
160
|
+
func = AsyncMock(side_effect=ServerDisconnectedError())
|
|
161
|
+
with patch("asyncio.sleep", new_callable=AsyncMock):
|
|
162
|
+
with raises(ServerDisconnectedError):
|
|
163
|
+
await specimen.retry_request(func, max_retries=2)
|
|
164
|
+
assert func.call_count == 2
|
|
165
|
+
|
|
166
|
+
@mark.asyncio
|
|
167
|
+
async def test_retry_request_client_os_error(self):
|
|
168
|
+
func = AsyncMock(
|
|
169
|
+
side_effect=ClientOSError(Mock(), OSError("OS error"))
|
|
170
|
+
)
|
|
171
|
+
with patch("asyncio.sleep", new_callable=AsyncMock):
|
|
172
|
+
with raises(ClientOSError):
|
|
173
|
+
await specimen.retry_request(func, max_retries=2)
|
|
174
|
+
assert func.call_count == 2
|
|
175
|
+
|
|
176
|
+
@mark.asyncio
|
|
177
|
+
async def test_retry_request_timeout_error(self):
|
|
178
|
+
func = AsyncMock(side_effect=asyncio.TimeoutError())
|
|
179
|
+
with patch("asyncio.sleep", new_callable=AsyncMock):
|
|
180
|
+
with raises(asyncio.TimeoutError):
|
|
181
|
+
await specimen.retry_request(func, max_retries=2)
|
|
182
|
+
assert func.call_count == 2
|
|
183
|
+
|
|
184
|
+
@mark.asyncio
|
|
185
|
+
async def test_retry_request_exponential_backoff(self):
|
|
186
|
+
func = AsyncMock(side_effect=[ServerTimeoutError(), "success"])
|
|
187
|
+
sleep_mock = AsyncMock()
|
|
188
|
+
with patch("asyncio.sleep", sleep_mock):
|
|
189
|
+
await specimen.retry_request(func, max_retries=3, retry_backoff=2)
|
|
190
|
+
# Should sleep with exponential backoff: 2^1 = 2 seconds
|
|
191
|
+
assert sleep_mock.call_count == 1
|
|
192
|
+
sleep_mock.assert_called_with(2)
|
|
193
|
+
|
|
194
|
+
|
|
195
|
+
class PingEndpointTests:
|
|
196
|
+
"""Tests for ping_endpoint (lines 72-140 in transfer.py)."""
|
|
197
|
+
|
|
198
|
+
def _make_ping_session_mock(self, response_status, response_text=""):
|
|
199
|
+
"""Build mocked ClientSession + session.get returning response_status."""
|
|
200
|
+
mock_resp = AsyncMock()
|
|
201
|
+
mock_resp.status = response_status
|
|
202
|
+
mock_resp.request_info = Mock()
|
|
203
|
+
mock_resp.history = ()
|
|
204
|
+
mock_resp.text = AsyncMock(return_value=response_text)
|
|
205
|
+
mock_resp_ctx = AsyncMock()
|
|
206
|
+
mock_resp_ctx.__aenter__ = AsyncMock(return_value=mock_resp)
|
|
207
|
+
mock_resp_ctx.__aexit__ = AsyncMock(return_value=None)
|
|
208
|
+
|
|
209
|
+
mock_session_instance = AsyncMock()
|
|
210
|
+
mock_session_instance.get = Mock(return_value=mock_resp_ctx)
|
|
211
|
+
|
|
212
|
+
mock_session_ctx = AsyncMock()
|
|
213
|
+
mock_session_ctx.__aenter__ = AsyncMock(
|
|
214
|
+
return_value=mock_session_instance
|
|
215
|
+
)
|
|
216
|
+
mock_session_ctx.__aexit__ = AsyncMock(return_value=None)
|
|
217
|
+
return mock_session_ctx
|
|
218
|
+
|
|
219
|
+
@mark.asyncio
|
|
220
|
+
async def test_ping_endpoint_success(self):
|
|
221
|
+
with patch(
|
|
222
|
+
"proximl.utils.transfer.aiohttp.ClientSession"
|
|
223
|
+
) as mock_session_class:
|
|
224
|
+
mock_session_ctx = self._make_ping_session_mock(200)
|
|
225
|
+
mock_session_class.return_value = mock_session_ctx
|
|
226
|
+
with patch("asyncio.sleep", new_callable=AsyncMock):
|
|
227
|
+
await specimen.ping_endpoint("https://host", "token")
|
|
228
|
+
mock_session_class.assert_called_once()
|
|
229
|
+
|
|
230
|
+
@mark.asyncio
|
|
231
|
+
async def test_ping_endpoint_http_error_retry_then_success(self):
|
|
232
|
+
call_count = [0]
|
|
233
|
+
|
|
234
|
+
def get_ctx(*args, **kwargs):
|
|
235
|
+
call_count[0] += 1
|
|
236
|
+
status = 200 if call_count[0] > 1 else 404
|
|
237
|
+
mock_resp = AsyncMock()
|
|
238
|
+
mock_resp.status = status
|
|
239
|
+
mock_resp.request_info = Mock()
|
|
240
|
+
mock_resp.history = ()
|
|
241
|
+
mock_resp.text = AsyncMock(return_value="Not Found")
|
|
242
|
+
mock_resp_ctx = AsyncMock()
|
|
243
|
+
mock_resp_ctx.__aenter__ = AsyncMock(return_value=mock_resp)
|
|
244
|
+
mock_resp_ctx.__aexit__ = AsyncMock(return_value=None)
|
|
245
|
+
return mock_resp_ctx
|
|
246
|
+
|
|
247
|
+
with patch(
|
|
248
|
+
"proximl.utils.transfer.aiohttp.ClientSession"
|
|
249
|
+
) as mock_session_class:
|
|
250
|
+
mock_session_instance = AsyncMock()
|
|
251
|
+
mock_session_instance.get = Mock(side_effect=get_ctx)
|
|
252
|
+
mock_session_ctx = AsyncMock()
|
|
253
|
+
mock_session_ctx.__aenter__ = AsyncMock(
|
|
254
|
+
return_value=mock_session_instance
|
|
255
|
+
)
|
|
256
|
+
mock_session_ctx.__aexit__ = AsyncMock(return_value=None)
|
|
257
|
+
mock_session_class.return_value = mock_session_ctx
|
|
258
|
+
with patch("asyncio.sleep", new_callable=AsyncMock):
|
|
259
|
+
await specimen.ping_endpoint(
|
|
260
|
+
"https://host", "token", max_retries=3
|
|
261
|
+
)
|
|
262
|
+
assert call_count[0] == 2
|
|
263
|
+
|
|
264
|
+
@mark.asyncio
|
|
265
|
+
async def test_ping_endpoint_http_error_max_retries(self):
|
|
266
|
+
mock_session_ctx = self._make_ping_session_mock(404, "Not Found")
|
|
267
|
+
with patch(
|
|
268
|
+
"proximl.utils.transfer.aiohttp.ClientSession"
|
|
269
|
+
) as mock_session_class:
|
|
270
|
+
mock_session_class.return_value = mock_session_ctx
|
|
271
|
+
with patch("asyncio.sleep", new_callable=AsyncMock):
|
|
272
|
+
with raises(
|
|
273
|
+
ConnectionError, match="ping failed after 2 attempts"
|
|
274
|
+
):
|
|
275
|
+
await specimen.ping_endpoint(
|
|
276
|
+
"https://host", "token", max_retries=2
|
|
277
|
+
)
|
|
278
|
+
|
|
279
|
+
@mark.asyncio
|
|
280
|
+
async def test_ping_endpoint_connector_error_dns_retries_then_success(
|
|
281
|
+
self,
|
|
282
|
+
):
|
|
283
|
+
call_count = [0]
|
|
284
|
+
|
|
285
|
+
def session_get(*args, **kwargs):
|
|
286
|
+
call_count[0] += 1
|
|
287
|
+
if call_count[0] == 1:
|
|
288
|
+
raise ClientConnectorError(
|
|
289
|
+
connection_key=Mock(), os_error=OSError("dns")
|
|
290
|
+
)
|
|
291
|
+
mock_resp = AsyncMock()
|
|
292
|
+
mock_resp.status = 200
|
|
293
|
+
mock_resp.request_info = Mock()
|
|
294
|
+
mock_resp.history = ()
|
|
295
|
+
mock_resp.text = AsyncMock(return_value="")
|
|
296
|
+
mock_resp_ctx = AsyncMock()
|
|
297
|
+
mock_resp_ctx.__aenter__ = AsyncMock(return_value=mock_resp)
|
|
298
|
+
mock_resp_ctx.__aexit__ = AsyncMock(return_value=None)
|
|
299
|
+
return mock_resp_ctx
|
|
300
|
+
|
|
301
|
+
with patch(
|
|
302
|
+
"proximl.utils.transfer.aiohttp.ClientSession"
|
|
303
|
+
) as mock_session_class:
|
|
304
|
+
mock_session_instance = AsyncMock()
|
|
305
|
+
mock_session_instance.get = Mock(side_effect=session_get)
|
|
306
|
+
mock_session_ctx = AsyncMock()
|
|
307
|
+
mock_session_ctx.__aenter__ = AsyncMock(
|
|
308
|
+
return_value=mock_session_instance
|
|
309
|
+
)
|
|
310
|
+
mock_session_ctx.__aexit__ = AsyncMock(return_value=None)
|
|
311
|
+
mock_session_class.return_value = mock_session_ctx
|
|
312
|
+
sleep_mock = AsyncMock()
|
|
313
|
+
with patch("asyncio.sleep", sleep_mock):
|
|
314
|
+
await specimen.ping_endpoint(
|
|
315
|
+
"https://host", "token", max_retries=2
|
|
316
|
+
)
|
|
317
|
+
assert call_count[0] == 2
|
|
318
|
+
sleep_mock.assert_called_once_with(specimen.DNS_INITIAL_DELAY)
|
|
319
|
+
|
|
320
|
+
@mark.asyncio
|
|
321
|
+
async def test_ping_endpoint_connector_error_max_retries(self):
|
|
322
|
+
with patch(
|
|
323
|
+
"proximl.utils.transfer.aiohttp.ClientSession"
|
|
324
|
+
) as mock_session_class:
|
|
325
|
+
mock_session_instance = AsyncMock()
|
|
326
|
+
mock_session_instance.get = Mock(
|
|
327
|
+
side_effect=ClientConnectorError(
|
|
328
|
+
connection_key=Mock(), os_error=OSError("dns")
|
|
329
|
+
)
|
|
330
|
+
)
|
|
331
|
+
mock_session_ctx = AsyncMock()
|
|
332
|
+
mock_session_ctx.__aenter__ = AsyncMock(
|
|
333
|
+
return_value=mock_session_instance
|
|
334
|
+
)
|
|
335
|
+
mock_session_ctx.__aexit__ = AsyncMock(return_value=None)
|
|
336
|
+
mock_session_class.return_value = mock_session_ctx
|
|
337
|
+
sleep_mock = AsyncMock()
|
|
338
|
+
with patch("asyncio.sleep", sleep_mock):
|
|
339
|
+
with raises(
|
|
340
|
+
ConnectionError,
|
|
341
|
+
match="ping failed after .* attempts due to DNS/connection",
|
|
342
|
+
):
|
|
343
|
+
await specimen.ping_endpoint(
|
|
344
|
+
"https://host", "token", max_retries=2
|
|
345
|
+
)
|
|
346
|
+
assert sleep_mock.call_count >= 1
|
|
347
|
+
|
|
348
|
+
@mark.asyncio
|
|
349
|
+
async def test_ping_endpoint_connector_error_backoff_after_first_retry(
|
|
350
|
+
self,
|
|
351
|
+
):
|
|
352
|
+
call_count = [0]
|
|
353
|
+
|
|
354
|
+
def session_get(*args, **kwargs):
|
|
355
|
+
call_count[0] += 1
|
|
356
|
+
if call_count[0] < 3:
|
|
357
|
+
raise ClientConnectorError(
|
|
358
|
+
connection_key=Mock(), os_error=OSError("dns")
|
|
359
|
+
)
|
|
360
|
+
mock_resp = AsyncMock()
|
|
361
|
+
mock_resp.status = 200
|
|
362
|
+
mock_resp.request_info = Mock()
|
|
363
|
+
mock_resp.history = ()
|
|
364
|
+
mock_resp.text = AsyncMock(return_value="")
|
|
365
|
+
mock_resp_ctx = AsyncMock()
|
|
366
|
+
mock_resp_ctx.__aenter__ = AsyncMock(return_value=mock_resp)
|
|
367
|
+
mock_resp_ctx.__aexit__ = AsyncMock(return_value=None)
|
|
368
|
+
return mock_resp_ctx
|
|
369
|
+
|
|
370
|
+
with patch(
|
|
371
|
+
"proximl.utils.transfer.aiohttp.ClientSession"
|
|
372
|
+
) as mock_session_class:
|
|
373
|
+
mock_session_instance = AsyncMock()
|
|
374
|
+
mock_session_instance.get = Mock(side_effect=session_get)
|
|
375
|
+
mock_session_ctx = AsyncMock()
|
|
376
|
+
mock_session_ctx.__aenter__ = AsyncMock(
|
|
377
|
+
return_value=mock_session_instance
|
|
378
|
+
)
|
|
379
|
+
mock_session_ctx.__aexit__ = AsyncMock(return_value=None)
|
|
380
|
+
mock_session_class.return_value = mock_session_ctx
|
|
381
|
+
sleep_mock = AsyncMock()
|
|
382
|
+
with patch("asyncio.sleep", sleep_mock):
|
|
383
|
+
await specimen.ping_endpoint(
|
|
384
|
+
"https://host", "token", max_retries=5, retry_backoff=2
|
|
385
|
+
)
|
|
386
|
+
assert call_count[0] == 3
|
|
387
|
+
assert sleep_mock.call_count == 2
|
|
388
|
+
sleep_mock.assert_any_call(specimen.DNS_INITIAL_DELAY)
|
|
389
|
+
sleep_mock.assert_any_call(2)
|
|
390
|
+
|
|
391
|
+
@mark.asyncio
|
|
392
|
+
async def test_ping_endpoint_other_error_retry_then_success(self):
|
|
393
|
+
call_count = [0]
|
|
394
|
+
|
|
395
|
+
def session_get(*args, **kwargs):
|
|
396
|
+
call_count[0] += 1
|
|
397
|
+
if call_count[0] == 1:
|
|
398
|
+
raise ServerDisconnectedError("disconnected")
|
|
399
|
+
mock_resp = AsyncMock()
|
|
400
|
+
mock_resp.status = 200
|
|
401
|
+
mock_resp.request_info = Mock()
|
|
402
|
+
mock_resp.history = ()
|
|
403
|
+
mock_resp.text = AsyncMock(return_value="")
|
|
404
|
+
mock_resp_ctx = AsyncMock()
|
|
405
|
+
mock_resp_ctx.__aenter__ = AsyncMock(return_value=mock_resp)
|
|
406
|
+
mock_resp_ctx.__aexit__ = AsyncMock(return_value=None)
|
|
407
|
+
return mock_resp_ctx
|
|
408
|
+
|
|
409
|
+
with patch(
|
|
410
|
+
"proximl.utils.transfer.aiohttp.ClientSession"
|
|
411
|
+
) as mock_session_class:
|
|
412
|
+
mock_session_instance = AsyncMock()
|
|
413
|
+
mock_session_instance.get = Mock(side_effect=session_get)
|
|
414
|
+
mock_session_ctx = AsyncMock()
|
|
415
|
+
mock_session_ctx.__aenter__ = AsyncMock(
|
|
416
|
+
return_value=mock_session_instance
|
|
417
|
+
)
|
|
418
|
+
mock_session_ctx.__aexit__ = AsyncMock(return_value=None)
|
|
419
|
+
mock_session_class.return_value = mock_session_ctx
|
|
420
|
+
with patch("asyncio.sleep", new_callable=AsyncMock):
|
|
421
|
+
await specimen.ping_endpoint(
|
|
422
|
+
"https://host", "token", max_retries=3
|
|
423
|
+
)
|
|
424
|
+
assert call_count[0] == 2
|
|
425
|
+
|
|
426
|
+
@mark.asyncio
|
|
427
|
+
async def test_ping_endpoint_other_error_max_retries(self):
|
|
428
|
+
with patch(
|
|
429
|
+
"proximl.utils.transfer.aiohttp.ClientSession"
|
|
430
|
+
) as mock_session_class:
|
|
431
|
+
mock_session_instance = AsyncMock()
|
|
432
|
+
mock_session_instance.get = Mock(
|
|
433
|
+
side_effect=ServerDisconnectedError("disconnected")
|
|
434
|
+
)
|
|
435
|
+
mock_session_ctx = AsyncMock()
|
|
436
|
+
mock_session_ctx.__aenter__ = AsyncMock(
|
|
437
|
+
return_value=mock_session_instance
|
|
438
|
+
)
|
|
439
|
+
mock_session_ctx.__aexit__ = AsyncMock(return_value=None)
|
|
440
|
+
mock_session_class.return_value = mock_session_ctx
|
|
441
|
+
with patch("asyncio.sleep", new_callable=AsyncMock):
|
|
442
|
+
with raises(
|
|
443
|
+
ConnectionError,
|
|
444
|
+
match="ping failed after 2 attempts",
|
|
445
|
+
):
|
|
446
|
+
await specimen.ping_endpoint(
|
|
447
|
+
"https://host", "token", max_retries=2
|
|
448
|
+
)
|
|
449
|
+
|
|
450
|
+
@mark.asyncio
|
|
451
|
+
async def test_ping_endpoint_client_payload_error_retry_then_success(self):
|
|
452
|
+
call_count = [0]
|
|
453
|
+
|
|
454
|
+
def session_get(*args, **kwargs):
|
|
455
|
+
call_count[0] += 1
|
|
456
|
+
if call_count[0] == 1:
|
|
457
|
+
raise ClientPayloadError("payload error")
|
|
458
|
+
mock_resp = AsyncMock()
|
|
459
|
+
mock_resp.status = 200
|
|
460
|
+
mock_resp.request_info = Mock()
|
|
461
|
+
mock_resp.history = ()
|
|
462
|
+
mock_resp.text = AsyncMock(return_value="")
|
|
463
|
+
mock_resp_ctx = AsyncMock()
|
|
464
|
+
mock_resp_ctx.__aenter__ = AsyncMock(return_value=mock_resp)
|
|
465
|
+
mock_resp_ctx.__aexit__ = AsyncMock(return_value=None)
|
|
466
|
+
return mock_resp_ctx
|
|
467
|
+
|
|
468
|
+
with patch(
|
|
469
|
+
"proximl.utils.transfer.aiohttp.ClientSession"
|
|
470
|
+
) as mock_session_class:
|
|
471
|
+
mock_session_instance = AsyncMock()
|
|
472
|
+
mock_session_instance.get = Mock(side_effect=session_get)
|
|
473
|
+
mock_session_ctx = AsyncMock()
|
|
474
|
+
mock_session_ctx.__aenter__ = AsyncMock(
|
|
475
|
+
return_value=mock_session_instance
|
|
476
|
+
)
|
|
477
|
+
mock_session_ctx.__aexit__ = AsyncMock(return_value=None)
|
|
478
|
+
mock_session_class.return_value = mock_session_ctx
|
|
479
|
+
with patch("asyncio.sleep", new_callable=AsyncMock):
|
|
480
|
+
await specimen.ping_endpoint(
|
|
481
|
+
"https://host", "token", max_retries=3
|
|
482
|
+
)
|
|
483
|
+
assert call_count[0] == 2
|
|
484
|
+
|
|
485
|
+
@mark.asyncio
|
|
486
|
+
async def test_ping_endpoint_timeout_error_retry_then_success(self):
|
|
487
|
+
call_count = [0]
|
|
488
|
+
|
|
489
|
+
def session_get(*args, **kwargs):
|
|
490
|
+
call_count[0] += 1
|
|
491
|
+
if call_count[0] == 1:
|
|
492
|
+
raise asyncio.TimeoutError()
|
|
493
|
+
mock_resp = AsyncMock()
|
|
494
|
+
mock_resp.status = 200
|
|
495
|
+
mock_resp.request_info = Mock()
|
|
496
|
+
mock_resp.history = ()
|
|
497
|
+
mock_resp.text = AsyncMock(return_value="")
|
|
498
|
+
mock_resp_ctx = AsyncMock()
|
|
499
|
+
mock_resp_ctx.__aenter__ = AsyncMock(return_value=mock_resp)
|
|
500
|
+
mock_resp_ctx.__aexit__ = AsyncMock(return_value=None)
|
|
501
|
+
return mock_resp_ctx
|
|
502
|
+
|
|
503
|
+
with patch(
|
|
504
|
+
"proximl.utils.transfer.aiohttp.ClientSession"
|
|
505
|
+
) as mock_session_class:
|
|
506
|
+
mock_session_instance = AsyncMock()
|
|
507
|
+
mock_session_instance.get = Mock(side_effect=session_get)
|
|
508
|
+
mock_session_ctx = AsyncMock()
|
|
509
|
+
mock_session_ctx.__aenter__ = AsyncMock(
|
|
510
|
+
return_value=mock_session_instance
|
|
511
|
+
)
|
|
512
|
+
mock_session_ctx.__aexit__ = AsyncMock(return_value=None)
|
|
513
|
+
mock_session_class.return_value = mock_session_ctx
|
|
514
|
+
with patch("asyncio.sleep", new_callable=AsyncMock):
|
|
515
|
+
await specimen.ping_endpoint(
|
|
516
|
+
"https://host", "token", max_retries=3
|
|
517
|
+
)
|
|
518
|
+
assert call_count[0] == 2
|
|
519
|
+
|
|
520
|
+
@mark.asyncio
|
|
521
|
+
async def test_ping_endpoint_normalizes_endpoint(self):
|
|
522
|
+
with patch(
|
|
523
|
+
"proximl.utils.transfer.aiohttp.ClientSession"
|
|
524
|
+
) as mock_session_class:
|
|
525
|
+
mock_session_ctx = self._make_ping_session_mock(200)
|
|
526
|
+
mock_session_class.return_value = mock_session_ctx
|
|
527
|
+
with patch("asyncio.sleep", new_callable=AsyncMock):
|
|
528
|
+
await specimen.ping_endpoint("host.no.scheme", "token")
|
|
529
|
+
call_kw = mock_session_instance = (
|
|
530
|
+
mock_session_class.return_value.__aenter__.return_value
|
|
531
|
+
)
|
|
532
|
+
get_call = call_kw.get.call_args
|
|
533
|
+
assert get_call is not None
|
|
534
|
+
(url,) = get_call[0]
|
|
535
|
+
assert url == "https://host.no.scheme/ping"
|
|
536
|
+
|
|
537
|
+
@mark.asyncio
|
|
538
|
+
async def test_ping_endpoint_sends_bearer_auth(self):
|
|
539
|
+
with patch(
|
|
540
|
+
"proximl.utils.transfer.aiohttp.ClientSession"
|
|
541
|
+
) as mock_session_class:
|
|
542
|
+
mock_session_ctx = self._make_ping_session_mock(200)
|
|
543
|
+
mock_session_class.return_value = mock_session_ctx
|
|
544
|
+
with patch("asyncio.sleep", new_callable=AsyncMock):
|
|
545
|
+
await specimen.ping_endpoint("https://host", "my_token")
|
|
546
|
+
sess = mock_session_class.return_value.__aenter__.return_value
|
|
547
|
+
sess.get.assert_called_once()
|
|
548
|
+
call_kw = sess.get.call_args[1]
|
|
549
|
+
assert call_kw["headers"]["Authorization"] == "Bearer my_token"
|
|
550
|
+
|
|
551
|
+
|
|
552
|
+
class UploadChunkTests:
|
|
553
|
+
@mark.asyncio
|
|
554
|
+
async def test_upload_chunk_success(self):
|
|
555
|
+
# Mock session.put to return an async context manager
|
|
556
|
+
mock_response = AsyncMock()
|
|
557
|
+
mock_response.status = 200
|
|
558
|
+
mock_response.release = AsyncMock()
|
|
559
|
+
mock_response.request_info = Mock()
|
|
560
|
+
mock_response.history = ()
|
|
561
|
+
mock_response.__aenter__ = AsyncMock(return_value=mock_response)
|
|
562
|
+
mock_response.__aexit__ = AsyncMock(return_value=None)
|
|
563
|
+
|
|
564
|
+
session = AsyncMock()
|
|
565
|
+
# session.put() should return the response directly (not a coroutine)
|
|
566
|
+
session.put = Mock(return_value=mock_response)
|
|
567
|
+
|
|
568
|
+
# Mock retry_request to actually call and await the function passed to it
|
|
569
|
+
# The function (_upload) does: async with session.put() as response: ...
|
|
570
|
+
async def mock_retry(func, *args, **kwargs):
|
|
571
|
+
# func is _upload, which does async with session.put() as response:
|
|
572
|
+
return await func(*args, **kwargs)
|
|
573
|
+
|
|
574
|
+
with patch(
|
|
575
|
+
"proximl.utils.transfer.retry_request", new_callable=AsyncMock
|
|
576
|
+
) as mock_retry_patch:
|
|
577
|
+
mock_retry_patch.side_effect = mock_retry
|
|
578
|
+
await specimen.upload_chunk(
|
|
579
|
+
session,
|
|
580
|
+
"https://example.com",
|
|
581
|
+
"token",
|
|
582
|
+
100,
|
|
583
|
+
b"data",
|
|
584
|
+
0,
|
|
585
|
+
)
|
|
586
|
+
mock_response.release.assert_called_once()
|
|
587
|
+
|
|
588
|
+
@mark.asyncio
|
|
589
|
+
async def test_upload_chunk_content_range_header(self):
|
|
590
|
+
session = AsyncMock()
|
|
591
|
+
response = AsyncMock()
|
|
592
|
+
response.status = 200
|
|
593
|
+
response.release = AsyncMock()
|
|
594
|
+
response.request_info = Mock()
|
|
595
|
+
response.history = ()
|
|
596
|
+
session.put = AsyncMock(return_value=response.__aenter__())
|
|
597
|
+
response.__aenter__ = AsyncMock(return_value=response)
|
|
598
|
+
response.__aexit__ = AsyncMock(return_value=None)
|
|
599
|
+
|
|
600
|
+
# Mock the response that _upload returns
|
|
601
|
+
mock_response_context = AsyncMock()
|
|
602
|
+
mock_response_context.status = 200
|
|
603
|
+
mock_response_context.release = AsyncMock()
|
|
604
|
+
mock_response_context.request_info = Mock()
|
|
605
|
+
mock_response_context.history = ()
|
|
606
|
+
mock_response_context.__aenter__ = AsyncMock(
|
|
607
|
+
return_value=mock_response_context
|
|
608
|
+
)
|
|
609
|
+
mock_response_context.__aexit__ = AsyncMock(return_value=None)
|
|
610
|
+
|
|
611
|
+
async def _upload(*args, **kwargs):
|
|
612
|
+
return mock_response_context
|
|
613
|
+
|
|
614
|
+
with patch("proximl.utils.transfer.retry_request") as mock_retry:
|
|
615
|
+
mock_retry.side_effect = _upload
|
|
616
|
+
await specimen.upload_chunk(
|
|
617
|
+
session,
|
|
618
|
+
"https://example.com",
|
|
619
|
+
"token",
|
|
620
|
+
100,
|
|
621
|
+
b"data",
|
|
622
|
+
10,
|
|
623
|
+
)
|
|
624
|
+
# Verify headers were set correctly by checking session.put was called
|
|
625
|
+
# Note: session.put is called inside _upload, but we're mocking retry_request
|
|
626
|
+
# So we verify the function completed successfully
|
|
627
|
+
assert True # Test passes if no exception
|
|
628
|
+
|
|
629
|
+
@mark.asyncio
|
|
630
|
+
async def test_upload_chunk_retry_status(self):
|
|
631
|
+
session = AsyncMock()
|
|
632
|
+
response = AsyncMock()
|
|
633
|
+
response.status = 502
|
|
634
|
+
response.text = AsyncMock(return_value="Bad Gateway")
|
|
635
|
+
response.request_info = Mock()
|
|
636
|
+
response.history = ()
|
|
637
|
+
session.put = AsyncMock(return_value=response.__aenter__())
|
|
638
|
+
response.__aenter__ = AsyncMock(return_value=response)
|
|
639
|
+
response.__aexit__ = AsyncMock(return_value=None)
|
|
640
|
+
|
|
641
|
+
with patch("proximl.utils.transfer.retry_request") as mock_retry:
|
|
642
|
+
|
|
643
|
+
async def _upload():
|
|
644
|
+
async with session.put(
|
|
645
|
+
f"https://example.com/upload",
|
|
646
|
+
headers={},
|
|
647
|
+
data=b"data",
|
|
648
|
+
timeout=30,
|
|
649
|
+
) as resp:
|
|
650
|
+
if resp.status == 502:
|
|
651
|
+
text = await resp.text()
|
|
652
|
+
raise ClientResponseError(
|
|
653
|
+
request_info=resp.request_info,
|
|
654
|
+
history=resp.history,
|
|
655
|
+
status=resp.status,
|
|
656
|
+
message=text,
|
|
657
|
+
)
|
|
658
|
+
|
|
659
|
+
mock_retry.side_effect = ClientResponseError(
|
|
660
|
+
request_info=Mock(),
|
|
661
|
+
history=(),
|
|
662
|
+
status=502,
|
|
663
|
+
message="Bad Gateway",
|
|
664
|
+
)
|
|
665
|
+
with raises(ClientResponseError):
|
|
666
|
+
await specimen.upload_chunk(
|
|
667
|
+
session,
|
|
668
|
+
"https://example.com",
|
|
669
|
+
"token",
|
|
670
|
+
100,
|
|
671
|
+
b"data",
|
|
672
|
+
0,
|
|
673
|
+
)
|
|
674
|
+
|
|
675
|
+
@mark.asyncio
|
|
676
|
+
async def test_upload_chunk_error_status(self):
|
|
677
|
+
session = AsyncMock()
|
|
678
|
+
response = AsyncMock()
|
|
679
|
+
response.status = 400
|
|
680
|
+
response.text = AsyncMock(return_value="Bad Request")
|
|
681
|
+
response.request_info = Mock()
|
|
682
|
+
response.history = ()
|
|
683
|
+
session.put = AsyncMock(return_value=response.__aenter__())
|
|
684
|
+
response.__aenter__ = AsyncMock(return_value=response)
|
|
685
|
+
response.__aexit__ = AsyncMock(return_value=None)
|
|
686
|
+
|
|
687
|
+
with patch("proximl.utils.transfer.retry_request") as mock_retry:
|
|
688
|
+
mock_retry.side_effect = ConnectionError(
|
|
689
|
+
"Chunk 0-3 failed with status 400: Bad Request"
|
|
690
|
+
)
|
|
691
|
+
with raises(ConnectionError, match="Chunk.*failed"):
|
|
692
|
+
await specimen.upload_chunk(
|
|
693
|
+
session,
|
|
694
|
+
"https://example.com",
|
|
695
|
+
"token",
|
|
696
|
+
100,
|
|
697
|
+
b"data",
|
|
698
|
+
0,
|
|
699
|
+
)
|
|
700
|
+
|
|
701
|
+
@mark.asyncio
|
|
702
|
+
async def test_upload_chunk_retry_status_direct(self):
|
|
703
|
+
"""Test upload_chunk retry status (lines 106-113) - direct execution"""
|
|
704
|
+
session = AsyncMock()
|
|
705
|
+
|
|
706
|
+
# Create a response with retry status (502)
|
|
707
|
+
mock_response = AsyncMock()
|
|
708
|
+
mock_response.status = 502
|
|
709
|
+
mock_response.text = AsyncMock(return_value="Bad Gateway")
|
|
710
|
+
mock_response.request_info = Mock()
|
|
711
|
+
mock_response.history = ()
|
|
712
|
+
mock_response.release = AsyncMock()
|
|
713
|
+
|
|
714
|
+
# Make session.put return an async context manager
|
|
715
|
+
class AwaitableContextManager:
|
|
716
|
+
def __init__(self, return_value):
|
|
717
|
+
self.return_value = return_value
|
|
718
|
+
|
|
719
|
+
def __await__(self):
|
|
720
|
+
yield
|
|
721
|
+
return self
|
|
722
|
+
|
|
723
|
+
async def __aenter__(self):
|
|
724
|
+
return self.return_value
|
|
725
|
+
|
|
726
|
+
async def __aexit__(self, *args):
|
|
727
|
+
return None
|
|
728
|
+
|
|
729
|
+
mock_put_context = AwaitableContextManager(mock_response)
|
|
730
|
+
session.put = Mock(return_value=mock_put_context)
|
|
731
|
+
|
|
732
|
+
# Mock retry_request to actually call the function passed to it
|
|
733
|
+
async def mock_retry(func, *args, **kwargs):
|
|
734
|
+
return await func(*args, **kwargs)
|
|
735
|
+
|
|
736
|
+
with patch(
|
|
737
|
+
"proximl.utils.transfer.retry_request", side_effect=mock_retry
|
|
738
|
+
):
|
|
739
|
+
with raises(ClientResponseError):
|
|
740
|
+
await specimen.upload_chunk(
|
|
741
|
+
session,
|
|
742
|
+
"https://example.com",
|
|
743
|
+
"token",
|
|
744
|
+
100,
|
|
745
|
+
b"data",
|
|
746
|
+
0,
|
|
747
|
+
)
|
|
748
|
+
|
|
749
|
+
@mark.asyncio
|
|
750
|
+
async def test_upload_chunk_error_status_direct(self):
|
|
751
|
+
"""Test upload_chunk error status (lines 114-118) - direct execution"""
|
|
752
|
+
session = AsyncMock()
|
|
753
|
+
|
|
754
|
+
# Create a response with non-retry error status (400)
|
|
755
|
+
mock_response = AsyncMock()
|
|
756
|
+
mock_response.status = 400
|
|
757
|
+
mock_response.text = AsyncMock(return_value="Bad Request")
|
|
758
|
+
mock_response.request_info = Mock()
|
|
759
|
+
mock_response.history = ()
|
|
760
|
+
mock_response.release = AsyncMock()
|
|
761
|
+
|
|
762
|
+
# Make session.put return an async context manager
|
|
763
|
+
class AwaitableContextManager:
|
|
764
|
+
def __init__(self, return_value):
|
|
765
|
+
self.return_value = return_value
|
|
766
|
+
|
|
767
|
+
def __await__(self):
|
|
768
|
+
yield
|
|
769
|
+
return self
|
|
770
|
+
|
|
771
|
+
async def __aenter__(self):
|
|
772
|
+
return self.return_value
|
|
773
|
+
|
|
774
|
+
async def __aexit__(self, *args):
|
|
775
|
+
return None
|
|
776
|
+
|
|
777
|
+
mock_put_context = AwaitableContextManager(mock_response)
|
|
778
|
+
session.put = Mock(return_value=mock_put_context)
|
|
779
|
+
|
|
780
|
+
# Mock retry_request to actually call the function passed to it
|
|
781
|
+
async def mock_retry(func, *args, **kwargs):
|
|
782
|
+
return await func(*args, **kwargs)
|
|
783
|
+
|
|
784
|
+
with patch(
|
|
785
|
+
"proximl.utils.transfer.retry_request", side_effect=mock_retry
|
|
786
|
+
):
|
|
787
|
+
with raises(
|
|
788
|
+
ConnectionError, match="Chunk.*failed with status 400"
|
|
789
|
+
):
|
|
790
|
+
await specimen.upload_chunk(
|
|
791
|
+
session,
|
|
792
|
+
"https://example.com",
|
|
793
|
+
"token",
|
|
794
|
+
100,
|
|
795
|
+
b"data",
|
|
796
|
+
0,
|
|
797
|
+
)
|
|
798
|
+
|
|
799
|
+
|
|
800
|
+
class UploadTests:
|
|
801
|
+
@mark.asyncio
|
|
802
|
+
async def test_upload_file_not_found(self):
|
|
803
|
+
with patch(
|
|
804
|
+
"proximl.utils.transfer.ping_endpoint", new_callable=AsyncMock
|
|
805
|
+
):
|
|
806
|
+
with raises(ValueError, match="Path not found"):
|
|
807
|
+
await specimen.upload(
|
|
808
|
+
"https://example.com", "token", "/nonexistent/path"
|
|
809
|
+
)
|
|
810
|
+
|
|
811
|
+
@mark.asyncio
|
|
812
|
+
async def test_upload_invalid_path_type(self):
|
|
813
|
+
# Test path that is neither file nor directory
|
|
814
|
+
# This is hard to create in practice, but we can mock it
|
|
815
|
+
with tempfile.NamedTemporaryFile() as tmp:
|
|
816
|
+
with patch(
|
|
817
|
+
"proximl.utils.transfer.ping_endpoint", new_callable=AsyncMock
|
|
818
|
+
):
|
|
819
|
+
with patch("os.path.isfile", return_value=False):
|
|
820
|
+
with patch("os.path.isdir", return_value=False):
|
|
821
|
+
with patch("os.path.exists", return_value=True):
|
|
822
|
+
with raises(
|
|
823
|
+
ValueError,
|
|
824
|
+
match="Path is neither a file nor directory",
|
|
825
|
+
):
|
|
826
|
+
await specimen.upload(
|
|
827
|
+
"example.com", "token", tmp.name
|
|
828
|
+
)
|
|
829
|
+
|
|
830
|
+
@mark.asyncio
|
|
831
|
+
async def test_upload_file(self):
|
|
832
|
+
with tempfile.NamedTemporaryFile() as tmp:
|
|
833
|
+
tmp.write(b"test content")
|
|
834
|
+
tmp.flush()
|
|
835
|
+
|
|
836
|
+
with patch(
|
|
837
|
+
"proximl.utils.transfer.ping_endpoint", new_callable=AsyncMock
|
|
838
|
+
):
|
|
839
|
+
with patch(
|
|
840
|
+
"asyncio.create_subprocess_exec"
|
|
841
|
+
) as mock_subprocess:
|
|
842
|
+
mock_process = AsyncMock()
|
|
843
|
+
mock_process.stdout.read = AsyncMock(
|
|
844
|
+
side_effect=[b"data", b""]
|
|
845
|
+
)
|
|
846
|
+
mock_process.returncode = 0
|
|
847
|
+
mock_process.wait = AsyncMock(return_value=0)
|
|
848
|
+
mock_process.stderr.read = AsyncMock(return_value=b"")
|
|
849
|
+
mock_subprocess.return_value = mock_process
|
|
850
|
+
|
|
851
|
+
with patch("aiohttp.ClientSession") as mock_session:
|
|
852
|
+
mock_session_instance = AsyncMock()
|
|
853
|
+
mock_session.return_value.__aenter__ = AsyncMock(
|
|
854
|
+
return_value=mock_session_instance
|
|
855
|
+
)
|
|
856
|
+
mock_session.return_value.__aexit__ = AsyncMock(
|
|
857
|
+
return_value=None
|
|
858
|
+
)
|
|
859
|
+
|
|
860
|
+
with patch(
|
|
861
|
+
"proximl.utils.transfer.upload_chunk",
|
|
862
|
+
new_callable=AsyncMock,
|
|
863
|
+
) as mock_upload_chunk:
|
|
864
|
+
mock_finalize_response = AsyncMock()
|
|
865
|
+
mock_finalize_response.status = 200
|
|
866
|
+
mock_finalize_response.json = AsyncMock(
|
|
867
|
+
return_value={"status": "ok"}
|
|
868
|
+
)
|
|
869
|
+
mock_finalize_response.__aenter__ = AsyncMock(
|
|
870
|
+
return_value=mock_finalize_response
|
|
871
|
+
)
|
|
872
|
+
mock_finalize_response.__aexit__ = AsyncMock(
|
|
873
|
+
return_value=None
|
|
874
|
+
)
|
|
875
|
+
|
|
876
|
+
# session.post() should return something that is both awaitable and an async context manager
|
|
877
|
+
class AwaitableContextManager:
|
|
878
|
+
def __init__(self, return_value):
|
|
879
|
+
self.return_value = return_value
|
|
880
|
+
|
|
881
|
+
def __await__(self):
|
|
882
|
+
yield
|
|
883
|
+
return self
|
|
884
|
+
|
|
885
|
+
async def __aenter__(self):
|
|
886
|
+
return self.return_value
|
|
887
|
+
|
|
888
|
+
async def __aexit__(self, *args):
|
|
889
|
+
return None
|
|
890
|
+
|
|
891
|
+
mock_post_context = AwaitableContextManager(
|
|
892
|
+
mock_finalize_response
|
|
893
|
+
)
|
|
894
|
+
mock_session_instance.post = Mock(
|
|
895
|
+
return_value=mock_post_context
|
|
896
|
+
)
|
|
897
|
+
|
|
898
|
+
await specimen.upload(
|
|
899
|
+
"example.com", "token", tmp.name
|
|
900
|
+
)
|
|
901
|
+
# Verify upload_chunk was called
|
|
902
|
+
assert mock_upload_chunk.called
|
|
903
|
+
|
|
904
|
+
@mark.asyncio
|
|
905
|
+
async def test_upload_directory(self):
|
|
906
|
+
with tempfile.TemporaryDirectory() as tmpdir:
|
|
907
|
+
test_file = os.path.join(tmpdir, "test.txt")
|
|
908
|
+
with open(test_file, "w") as f:
|
|
909
|
+
f.write("test content")
|
|
910
|
+
|
|
911
|
+
with patch(
|
|
912
|
+
"proximl.utils.transfer.ping_endpoint", new_callable=AsyncMock
|
|
913
|
+
):
|
|
914
|
+
with patch(
|
|
915
|
+
"asyncio.create_subprocess_exec"
|
|
916
|
+
) as mock_subprocess:
|
|
917
|
+
mock_process = AsyncMock()
|
|
918
|
+
mock_process.stdout.read = AsyncMock(
|
|
919
|
+
side_effect=[b"data", b""]
|
|
920
|
+
)
|
|
921
|
+
mock_process.returncode = 0
|
|
922
|
+
mock_process.wait = AsyncMock(return_value=0)
|
|
923
|
+
mock_process.stderr.read = AsyncMock(return_value=b"")
|
|
924
|
+
mock_subprocess.return_value = mock_process
|
|
925
|
+
|
|
926
|
+
with patch("aiohttp.ClientSession") as mock_session:
|
|
927
|
+
mock_session_instance = AsyncMock()
|
|
928
|
+
mock_session.return_value.__aenter__ = AsyncMock(
|
|
929
|
+
return_value=mock_session_instance
|
|
930
|
+
)
|
|
931
|
+
mock_session.return_value.__aexit__ = AsyncMock(
|
|
932
|
+
return_value=None
|
|
933
|
+
)
|
|
934
|
+
|
|
935
|
+
with patch(
|
|
936
|
+
"proximl.utils.transfer.upload_chunk",
|
|
937
|
+
new_callable=AsyncMock,
|
|
938
|
+
) as mock_upload_chunk:
|
|
939
|
+
mock_finalize_response = AsyncMock()
|
|
940
|
+
mock_finalize_response.status = 200
|
|
941
|
+
mock_finalize_response.json = AsyncMock(
|
|
942
|
+
return_value={"status": "ok"}
|
|
943
|
+
)
|
|
944
|
+
mock_finalize_response.__aenter__ = AsyncMock(
|
|
945
|
+
return_value=mock_finalize_response
|
|
946
|
+
)
|
|
947
|
+
mock_finalize_response.__aexit__ = AsyncMock(
|
|
948
|
+
return_value=None
|
|
949
|
+
)
|
|
950
|
+
|
|
951
|
+
# session.post() should return something that is both awaitable and an async context manager
|
|
952
|
+
class AwaitableContextManager:
|
|
953
|
+
def __init__(self, return_value):
|
|
954
|
+
self.return_value = return_value
|
|
955
|
+
|
|
956
|
+
def __await__(self):
|
|
957
|
+
yield
|
|
958
|
+
return self
|
|
959
|
+
|
|
960
|
+
async def __aenter__(self):
|
|
961
|
+
return self.return_value
|
|
962
|
+
|
|
963
|
+
async def __aexit__(self, *args):
|
|
964
|
+
return None
|
|
965
|
+
|
|
966
|
+
mock_post_context = AwaitableContextManager(
|
|
967
|
+
mock_finalize_response
|
|
968
|
+
)
|
|
969
|
+
mock_session_instance.post = Mock(
|
|
970
|
+
return_value=mock_post_context
|
|
971
|
+
)
|
|
972
|
+
|
|
973
|
+
await specimen.upload(
|
|
974
|
+
"example.com", "token", tmpdir
|
|
975
|
+
)
|
|
976
|
+
# Verify upload_chunk was called
|
|
977
|
+
assert mock_upload_chunk.called
|
|
978
|
+
|
|
979
|
+
@mark.asyncio
|
|
980
|
+
async def test_upload_tar_command_failure(self):
|
|
981
|
+
with tempfile.NamedTemporaryFile() as tmp:
|
|
982
|
+
with patch("asyncio.create_subprocess_exec") as mock_subprocess:
|
|
983
|
+
mock_process = AsyncMock()
|
|
984
|
+
mock_process.stdout.read = AsyncMock(
|
|
985
|
+
side_effect=[b"data", b""]
|
|
986
|
+
)
|
|
987
|
+
mock_process.wait = AsyncMock(return_value=1)
|
|
988
|
+
mock_process.stderr.read = AsyncMock(return_value=b"tar error")
|
|
989
|
+
mock_subprocess.return_value = mock_process
|
|
990
|
+
|
|
991
|
+
with patch("aiohttp.ClientSession") as mock_session:
|
|
992
|
+
mock_session_instance = AsyncMock()
|
|
993
|
+
mock_session.return_value.__aenter__ = AsyncMock(
|
|
994
|
+
return_value=mock_session_instance
|
|
995
|
+
)
|
|
996
|
+
mock_session.return_value.__aexit__ = AsyncMock(
|
|
997
|
+
return_value=None
|
|
998
|
+
)
|
|
999
|
+
with patch(
|
|
1000
|
+
"proximl.utils.transfer.ping_endpoint",
|
|
1001
|
+
new_callable=AsyncMock,
|
|
1002
|
+
):
|
|
1003
|
+
with patch(
|
|
1004
|
+
"proximl.utils.transfer.upload_chunk",
|
|
1005
|
+
new_callable=AsyncMock,
|
|
1006
|
+
):
|
|
1007
|
+
with raises(
|
|
1008
|
+
ProxiMLException, match="tar command failed"
|
|
1009
|
+
):
|
|
1010
|
+
await specimen.upload(
|
|
1011
|
+
"example.com", "token", tmp.name
|
|
1012
|
+
)
|
|
1013
|
+
|
|
1014
|
+
@mark.asyncio
|
|
1015
|
+
async def test_upload_tar_command_failure_no_stderr(self):
|
|
1016
|
+
with tempfile.NamedTemporaryFile() as tmp:
|
|
1017
|
+
with patch("asyncio.create_subprocess_exec") as mock_subprocess:
|
|
1018
|
+
mock_process = AsyncMock()
|
|
1019
|
+
mock_process.stdout.read = AsyncMock(
|
|
1020
|
+
side_effect=[b"data", b""]
|
|
1021
|
+
)
|
|
1022
|
+
mock_process.wait = AsyncMock(return_value=1)
|
|
1023
|
+
mock_process.stderr.read = AsyncMock(return_value=None)
|
|
1024
|
+
mock_subprocess.return_value = mock_process
|
|
1025
|
+
|
|
1026
|
+
with patch("aiohttp.ClientSession") as mock_session:
|
|
1027
|
+
mock_session_instance = AsyncMock()
|
|
1028
|
+
mock_session.return_value.__aenter__ = AsyncMock(
|
|
1029
|
+
return_value=mock_session_instance
|
|
1030
|
+
)
|
|
1031
|
+
mock_session.return_value.__aexit__ = AsyncMock(
|
|
1032
|
+
return_value=None
|
|
1033
|
+
)
|
|
1034
|
+
with patch(
|
|
1035
|
+
"proximl.utils.transfer.ping_endpoint",
|
|
1036
|
+
new_callable=AsyncMock,
|
|
1037
|
+
):
|
|
1038
|
+
with patch(
|
|
1039
|
+
"proximl.utils.transfer.upload_chunk",
|
|
1040
|
+
new_callable=AsyncMock,
|
|
1041
|
+
):
|
|
1042
|
+
with raises(
|
|
1043
|
+
ProxiMLException, match="tar command failed"
|
|
1044
|
+
):
|
|
1045
|
+
await specimen.upload(
|
|
1046
|
+
"example.com", "token", tmp.name
|
|
1047
|
+
)
|
|
1048
|
+
|
|
1049
|
+
@mark.asyncio
|
|
1050
|
+
async def test_upload_finalize_failure(self):
|
|
1051
|
+
with tempfile.NamedTemporaryFile() as tmp:
|
|
1052
|
+
with patch(
|
|
1053
|
+
"proximl.utils.transfer.ping_endpoint", new_callable=AsyncMock
|
|
1054
|
+
):
|
|
1055
|
+
with patch(
|
|
1056
|
+
"asyncio.create_subprocess_exec"
|
|
1057
|
+
) as mock_subprocess:
|
|
1058
|
+
mock_process = AsyncMock()
|
|
1059
|
+
mock_process.stdout.read = AsyncMock(
|
|
1060
|
+
side_effect=[b"data", b""]
|
|
1061
|
+
)
|
|
1062
|
+
mock_process.returncode = 0
|
|
1063
|
+
mock_process.wait = AsyncMock(return_value=0)
|
|
1064
|
+
mock_process.stderr.read = AsyncMock(return_value=b"")
|
|
1065
|
+
mock_subprocess.return_value = mock_process
|
|
1066
|
+
|
|
1067
|
+
with patch("aiohttp.ClientSession") as mock_session:
|
|
1068
|
+
mock_session_instance = AsyncMock()
|
|
1069
|
+
mock_session.return_value.__aenter__ = AsyncMock(
|
|
1070
|
+
return_value=mock_session_instance
|
|
1071
|
+
)
|
|
1072
|
+
mock_session.return_value.__aexit__ = AsyncMock(
|
|
1073
|
+
return_value=None
|
|
1074
|
+
)
|
|
1075
|
+
|
|
1076
|
+
with patch(
|
|
1077
|
+
"proximl.utils.transfer.upload_chunk",
|
|
1078
|
+
new_callable=AsyncMock,
|
|
1079
|
+
) as mock_upload_chunk:
|
|
1080
|
+
mock_finalize_response = AsyncMock()
|
|
1081
|
+
mock_finalize_response.status = 500
|
|
1082
|
+
mock_finalize_response.text = AsyncMock(
|
|
1083
|
+
return_value="Finalize error"
|
|
1084
|
+
)
|
|
1085
|
+
mock_finalize_response.__aenter__ = AsyncMock(
|
|
1086
|
+
return_value=mock_finalize_response
|
|
1087
|
+
)
|
|
1088
|
+
mock_finalize_response.__aexit__ = AsyncMock(
|
|
1089
|
+
return_value=None
|
|
1090
|
+
)
|
|
1091
|
+
|
|
1092
|
+
# session.post() should return something that is both awaitable and an async context manager
|
|
1093
|
+
class AwaitableContextManager:
|
|
1094
|
+
def __init__(self, return_value):
|
|
1095
|
+
self.return_value = return_value
|
|
1096
|
+
|
|
1097
|
+
def __await__(self):
|
|
1098
|
+
yield
|
|
1099
|
+
return self
|
|
1100
|
+
|
|
1101
|
+
async def __aenter__(self):
|
|
1102
|
+
return self.return_value
|
|
1103
|
+
|
|
1104
|
+
async def __aexit__(self, *args):
|
|
1105
|
+
return None
|
|
1106
|
+
|
|
1107
|
+
mock_post_context = AwaitableContextManager(
|
|
1108
|
+
mock_finalize_response
|
|
1109
|
+
)
|
|
1110
|
+
mock_session_instance.post = Mock(
|
|
1111
|
+
return_value=mock_post_context
|
|
1112
|
+
)
|
|
1113
|
+
|
|
1114
|
+
with patch(
|
|
1115
|
+
"proximl.utils.transfer.ping_endpoint",
|
|
1116
|
+
new_callable=AsyncMock,
|
|
1117
|
+
):
|
|
1118
|
+
with raises(
|
|
1119
|
+
ConnectionError, match="Finalize failed"
|
|
1120
|
+
):
|
|
1121
|
+
await specimen.upload(
|
|
1122
|
+
"example.com", "token", tmp.name
|
|
1123
|
+
)
|
|
1124
|
+
# Verify upload_chunk was called before finalize
|
|
1125
|
+
assert mock_upload_chunk.called
|
|
1126
|
+
|
|
1127
|
+
@mark.asyncio
|
|
1128
|
+
async def test_upload_multiple_chunks(self):
|
|
1129
|
+
with tempfile.NamedTemporaryFile() as tmp:
|
|
1130
|
+
tmp.write(b"x" * (10 * 1024 * 1024)) # 10MB file
|
|
1131
|
+
tmp.flush()
|
|
1132
|
+
|
|
1133
|
+
with patch("asyncio.create_subprocess_exec") as mock_subprocess:
|
|
1134
|
+
mock_process = AsyncMock()
|
|
1135
|
+
# Simulate multiple chunks
|
|
1136
|
+
mock_process.stdout.read = AsyncMock(
|
|
1137
|
+
side_effect=[
|
|
1138
|
+
b"x" * (5 * 1024 * 1024),
|
|
1139
|
+
b"x" * (5 * 1024 * 1024),
|
|
1140
|
+
b"",
|
|
1141
|
+
]
|
|
1142
|
+
)
|
|
1143
|
+
mock_process.returncode = 0
|
|
1144
|
+
mock_process.wait = AsyncMock(return_value=0)
|
|
1145
|
+
mock_process.stderr.read = AsyncMock(return_value=b"")
|
|
1146
|
+
mock_subprocess.return_value = mock_process
|
|
1147
|
+
|
|
1148
|
+
with patch("aiohttp.ClientSession") as mock_session:
|
|
1149
|
+
mock_session_instance = AsyncMock()
|
|
1150
|
+
mock_session.return_value.__aenter__ = AsyncMock(
|
|
1151
|
+
return_value=mock_session_instance
|
|
1152
|
+
)
|
|
1153
|
+
mock_session.return_value.__aexit__ = AsyncMock(
|
|
1154
|
+
return_value=None
|
|
1155
|
+
)
|
|
1156
|
+
|
|
1157
|
+
upload_chunk_mock = AsyncMock()
|
|
1158
|
+
with patch(
|
|
1159
|
+
"proximl.utils.transfer.upload_chunk",
|
|
1160
|
+
upload_chunk_mock,
|
|
1161
|
+
):
|
|
1162
|
+
mock_finalize_response = AsyncMock()
|
|
1163
|
+
mock_finalize_response.status = 200
|
|
1164
|
+
mock_finalize_response.json = AsyncMock(
|
|
1165
|
+
return_value={"status": "ok"}
|
|
1166
|
+
)
|
|
1167
|
+
mock_finalize_response.__aenter__ = AsyncMock(
|
|
1168
|
+
return_value=mock_finalize_response
|
|
1169
|
+
)
|
|
1170
|
+
mock_finalize_response.__aexit__ = AsyncMock(
|
|
1171
|
+
return_value=None
|
|
1172
|
+
)
|
|
1173
|
+
|
|
1174
|
+
# session.post() should return something that is both awaitable and an async context manager
|
|
1175
|
+
class AwaitableContextManager:
|
|
1176
|
+
def __init__(self, return_value):
|
|
1177
|
+
self.return_value = return_value
|
|
1178
|
+
|
|
1179
|
+
def __await__(self):
|
|
1180
|
+
yield
|
|
1181
|
+
return self
|
|
1182
|
+
|
|
1183
|
+
async def __aenter__(self):
|
|
1184
|
+
return self.return_value
|
|
1185
|
+
|
|
1186
|
+
async def __aexit__(self, *args):
|
|
1187
|
+
return None
|
|
1188
|
+
|
|
1189
|
+
mock_post_context = AwaitableContextManager(
|
|
1190
|
+
mock_finalize_response
|
|
1191
|
+
)
|
|
1192
|
+
mock_session_instance.post = Mock(
|
|
1193
|
+
return_value=mock_post_context
|
|
1194
|
+
)
|
|
1195
|
+
|
|
1196
|
+
with patch(
|
|
1197
|
+
"proximl.utils.transfer.ping_endpoint",
|
|
1198
|
+
new_callable=AsyncMock,
|
|
1199
|
+
):
|
|
1200
|
+
await specimen.upload(
|
|
1201
|
+
"example.com", "token", tmp.name
|
|
1202
|
+
)
|
|
1203
|
+
# Should have called upload_chunk twice (one per chunk)
|
|
1204
|
+
assert upload_chunk_mock.call_count == 2
|
|
1205
|
+
|
|
1206
|
+
|
|
1207
|
+
class DownloadTests:
|
|
1208
|
+
@mark.asyncio
|
|
1209
|
+
async def test_download_creates_directory(self):
|
|
1210
|
+
with tempfile.TemporaryDirectory() as tmpdir:
|
|
1211
|
+
target_dir = os.path.join(tmpdir, "new_dir")
|
|
1212
|
+
|
|
1213
|
+
with patch("aiohttp.ClientSession") as mock_session:
|
|
1214
|
+
mock_session_instance = AsyncMock()
|
|
1215
|
+
mock_session.return_value.__aenter__ = AsyncMock(
|
|
1216
|
+
return_value=mock_session_instance
|
|
1217
|
+
)
|
|
1218
|
+
mock_session.return_value.__aexit__ = AsyncMock(
|
|
1219
|
+
return_value=None
|
|
1220
|
+
)
|
|
1221
|
+
|
|
1222
|
+
# Mock /info endpoint returning TAR mode
|
|
1223
|
+
def mock_get(*args, **kwargs):
|
|
1224
|
+
url = args[0] if args else kwargs.get("url", "")
|
|
1225
|
+
# Handle /ping endpoint for ping_endpoint
|
|
1226
|
+
if "/ping" in url:
|
|
1227
|
+
mock_resp = AsyncMock()
|
|
1228
|
+
mock_resp.status = 200
|
|
1229
|
+
mock_resp.__aenter__ = AsyncMock(
|
|
1230
|
+
return_value=mock_resp
|
|
1231
|
+
)
|
|
1232
|
+
mock_resp.__aexit__ = AsyncMock(return_value=None)
|
|
1233
|
+
return mock_resp
|
|
1234
|
+
# Handle /info endpoint
|
|
1235
|
+
elif "/info" in url:
|
|
1236
|
+
mock_resp = AsyncMock()
|
|
1237
|
+
mock_resp.status = 200
|
|
1238
|
+
mock_resp.json = AsyncMock(
|
|
1239
|
+
return_value={"archive": False}
|
|
1240
|
+
)
|
|
1241
|
+
mock_resp.text = AsyncMock(return_value="")
|
|
1242
|
+
mock_resp.__aenter__ = AsyncMock(
|
|
1243
|
+
return_value=mock_resp
|
|
1244
|
+
)
|
|
1245
|
+
mock_resp.__aexit__ = AsyncMock(return_value=None)
|
|
1246
|
+
return mock_resp
|
|
1247
|
+
else:
|
|
1248
|
+
# For /download endpoint, return awaitable that resolves to response
|
|
1249
|
+
async def get_download_response():
|
|
1250
|
+
mock_resp = AsyncMock()
|
|
1251
|
+
mock_resp.status = 200
|
|
1252
|
+
mock_resp.headers = {
|
|
1253
|
+
"Content-Type": "application/x-tar"
|
|
1254
|
+
}
|
|
1255
|
+
|
|
1256
|
+
async def chunk_iter():
|
|
1257
|
+
yield b"tar data"
|
|
1258
|
+
yield b""
|
|
1259
|
+
|
|
1260
|
+
mock_resp.content.iter_chunked = (
|
|
1261
|
+
lambda size: chunk_iter()
|
|
1262
|
+
)
|
|
1263
|
+
mock_resp.close = Mock()
|
|
1264
|
+
return mock_resp
|
|
1265
|
+
|
|
1266
|
+
return get_download_response()
|
|
1267
|
+
|
|
1268
|
+
mock_session_instance.get = Mock(side_effect=mock_get)
|
|
1269
|
+
|
|
1270
|
+
# retry_request is called 3 times: _get_info, _download, _finalize
|
|
1271
|
+
async def mock_retry(func, *args, **kwargs):
|
|
1272
|
+
return await func(*args, **kwargs)
|
|
1273
|
+
|
|
1274
|
+
with patch(
|
|
1275
|
+
"proximl.utils.transfer.retry_request",
|
|
1276
|
+
side_effect=mock_retry,
|
|
1277
|
+
):
|
|
1278
|
+
with patch(
|
|
1279
|
+
"asyncio.create_subprocess_exec"
|
|
1280
|
+
) as mock_subprocess:
|
|
1281
|
+
mock_process = AsyncMock()
|
|
1282
|
+
mock_process.stdin = Mock()
|
|
1283
|
+
mock_process.stdin.write = Mock()
|
|
1284
|
+
mock_process.stdin.drain = AsyncMock()
|
|
1285
|
+
mock_process.stdin.close = Mock()
|
|
1286
|
+
mock_process.returncode = 0
|
|
1287
|
+
mock_process.wait = AsyncMock(return_value=0)
|
|
1288
|
+
mock_process.stderr.read = AsyncMock(return_value=b"")
|
|
1289
|
+
mock_subprocess.return_value = mock_process
|
|
1290
|
+
|
|
1291
|
+
mock_finalize_response = AsyncMock()
|
|
1292
|
+
mock_finalize_response.status = 200
|
|
1293
|
+
mock_finalize_response.json = AsyncMock(
|
|
1294
|
+
return_value={"status": "ok"}
|
|
1295
|
+
)
|
|
1296
|
+
mock_finalize_response.__aenter__ = AsyncMock(
|
|
1297
|
+
return_value=mock_finalize_response
|
|
1298
|
+
)
|
|
1299
|
+
mock_finalize_response.__aexit__ = AsyncMock(
|
|
1300
|
+
return_value=None
|
|
1301
|
+
)
|
|
1302
|
+
|
|
1303
|
+
# session.post() should return something that is both awaitable and an async context manager
|
|
1304
|
+
class AwaitableContextManager:
|
|
1305
|
+
def __init__(self, return_value):
|
|
1306
|
+
self.return_value = return_value
|
|
1307
|
+
|
|
1308
|
+
def __await__(self):
|
|
1309
|
+
yield
|
|
1310
|
+
return self
|
|
1311
|
+
|
|
1312
|
+
async def __aenter__(self):
|
|
1313
|
+
return self.return_value
|
|
1314
|
+
|
|
1315
|
+
async def __aexit__(self, *args):
|
|
1316
|
+
return None
|
|
1317
|
+
|
|
1318
|
+
mock_post_context = AwaitableContextManager(
|
|
1319
|
+
mock_finalize_response
|
|
1320
|
+
)
|
|
1321
|
+
mock_session_instance.post = Mock(
|
|
1322
|
+
return_value=mock_post_context
|
|
1323
|
+
)
|
|
1324
|
+
|
|
1325
|
+
with patch(
|
|
1326
|
+
"proximl.utils.transfer.ping_endpoint",
|
|
1327
|
+
new_callable=AsyncMock,
|
|
1328
|
+
):
|
|
1329
|
+
await specimen.download(
|
|
1330
|
+
"example.com", "token", target_dir
|
|
1331
|
+
)
|
|
1332
|
+
|
|
1333
|
+
assert os.path.isdir(target_dir)
|
|
1334
|
+
|
|
1335
|
+
@mark.asyncio
|
|
1336
|
+
async def test_download_tar_mode(self):
|
|
1337
|
+
with tempfile.TemporaryDirectory() as tmpdir:
|
|
1338
|
+
with patch("aiohttp.ClientSession") as mock_session:
|
|
1339
|
+
mock_session_instance = AsyncMock()
|
|
1340
|
+
mock_session.return_value.__aenter__ = AsyncMock(
|
|
1341
|
+
return_value=mock_session_instance
|
|
1342
|
+
)
|
|
1343
|
+
mock_session.return_value.__aexit__ = AsyncMock(
|
|
1344
|
+
return_value=None
|
|
1345
|
+
)
|
|
1346
|
+
|
|
1347
|
+
# Mock /info endpoint - needs to return async context manager for async with
|
|
1348
|
+
def mock_get(*args, **kwargs):
|
|
1349
|
+
url = args[0] if args else kwargs.get("url", "")
|
|
1350
|
+
if "/info" in url or "/ping" in url:
|
|
1351
|
+
mock_resp = AsyncMock()
|
|
1352
|
+
mock_resp.status = 200
|
|
1353
|
+
mock_resp.json = AsyncMock(
|
|
1354
|
+
return_value={"archive": False}
|
|
1355
|
+
)
|
|
1356
|
+
mock_resp.text = AsyncMock(return_value="")
|
|
1357
|
+
# Configure as async context manager
|
|
1358
|
+
mock_resp.__aenter__ = AsyncMock(
|
|
1359
|
+
return_value=mock_resp
|
|
1360
|
+
)
|
|
1361
|
+
mock_resp.__aexit__ = AsyncMock(return_value=None)
|
|
1362
|
+
return mock_resp
|
|
1363
|
+
else:
|
|
1364
|
+
# For /download endpoint, return awaitable that resolves to response
|
|
1365
|
+
async def get_download_response():
|
|
1366
|
+
mock_resp = AsyncMock()
|
|
1367
|
+
mock_resp.status = 200
|
|
1368
|
+
mock_resp.headers = {
|
|
1369
|
+
"Content-Type": "application/x-tar"
|
|
1370
|
+
}
|
|
1371
|
+
|
|
1372
|
+
async def chunk_iter():
|
|
1373
|
+
yield b"tar data"
|
|
1374
|
+
yield b""
|
|
1375
|
+
|
|
1376
|
+
mock_resp.content.iter_chunked = (
|
|
1377
|
+
lambda size: chunk_iter()
|
|
1378
|
+
)
|
|
1379
|
+
mock_resp.close = Mock()
|
|
1380
|
+
return mock_resp
|
|
1381
|
+
|
|
1382
|
+
return get_download_response()
|
|
1383
|
+
|
|
1384
|
+
mock_session_instance.get = Mock(side_effect=mock_get)
|
|
1385
|
+
|
|
1386
|
+
# retry_request is called 3 times: _get_info, _download, _finalize
|
|
1387
|
+
async def mock_retry(func, *args, **kwargs):
|
|
1388
|
+
return await func(*args, **kwargs)
|
|
1389
|
+
|
|
1390
|
+
with patch(
|
|
1391
|
+
"proximl.utils.transfer.retry_request",
|
|
1392
|
+
side_effect=mock_retry,
|
|
1393
|
+
):
|
|
1394
|
+
with patch(
|
|
1395
|
+
"asyncio.create_subprocess_exec"
|
|
1396
|
+
) as mock_subprocess:
|
|
1397
|
+
mock_process = AsyncMock()
|
|
1398
|
+
mock_process.stdin = Mock()
|
|
1399
|
+
mock_process.stdin.write = Mock()
|
|
1400
|
+
mock_process.stdin.drain = AsyncMock()
|
|
1401
|
+
mock_process.stdin.close = Mock()
|
|
1402
|
+
mock_process.returncode = 0
|
|
1403
|
+
mock_process.wait = AsyncMock(return_value=0)
|
|
1404
|
+
mock_process.stderr.read = AsyncMock(return_value=b"")
|
|
1405
|
+
mock_subprocess.return_value = mock_process
|
|
1406
|
+
|
|
1407
|
+
mock_finalize_response = AsyncMock()
|
|
1408
|
+
mock_finalize_response.status = 200
|
|
1409
|
+
mock_finalize_response.json = AsyncMock(
|
|
1410
|
+
return_value={"status": "ok"}
|
|
1411
|
+
)
|
|
1412
|
+
mock_finalize_response.__aenter__ = AsyncMock(
|
|
1413
|
+
return_value=mock_finalize_response
|
|
1414
|
+
)
|
|
1415
|
+
mock_finalize_response.__aexit__ = AsyncMock(
|
|
1416
|
+
return_value=None
|
|
1417
|
+
)
|
|
1418
|
+
|
|
1419
|
+
# session.post() should return something that is both awaitable and an async context manager
|
|
1420
|
+
class AwaitableContextManager:
|
|
1421
|
+
def __init__(self, return_value):
|
|
1422
|
+
self.return_value = return_value
|
|
1423
|
+
|
|
1424
|
+
def __await__(self):
|
|
1425
|
+
yield
|
|
1426
|
+
return self
|
|
1427
|
+
|
|
1428
|
+
async def __aenter__(self):
|
|
1429
|
+
return self.return_value
|
|
1430
|
+
|
|
1431
|
+
async def __aexit__(self, *args):
|
|
1432
|
+
return None
|
|
1433
|
+
|
|
1434
|
+
mock_post_context = AwaitableContextManager(
|
|
1435
|
+
mock_finalize_response
|
|
1436
|
+
)
|
|
1437
|
+
mock_session_instance.post = Mock(
|
|
1438
|
+
return_value=mock_post_context
|
|
1439
|
+
)
|
|
1440
|
+
|
|
1441
|
+
with patch(
|
|
1442
|
+
"proximl.utils.transfer.ping_endpoint",
|
|
1443
|
+
new_callable=AsyncMock,
|
|
1444
|
+
):
|
|
1445
|
+
await specimen.download(
|
|
1446
|
+
"example.com", "token", tmpdir
|
|
1447
|
+
)
|
|
1448
|
+
|
|
1449
|
+
@mark.asyncio
|
|
1450
|
+
async def test_download_zip_mode(self):
|
|
1451
|
+
with tempfile.TemporaryDirectory() as tmpdir:
|
|
1452
|
+
with patch("aiohttp.ClientSession") as mock_session:
|
|
1453
|
+
mock_session_instance = AsyncMock()
|
|
1454
|
+
mock_session.return_value.__aenter__ = AsyncMock(
|
|
1455
|
+
return_value=mock_session_instance
|
|
1456
|
+
)
|
|
1457
|
+
mock_session.return_value.__aexit__ = AsyncMock(
|
|
1458
|
+
return_value=None
|
|
1459
|
+
)
|
|
1460
|
+
|
|
1461
|
+
# Mock /info endpoint returning ZIP mode - needs to return async context manager for async with
|
|
1462
|
+
def mock_get(*args, **kwargs):
|
|
1463
|
+
url = args[0] if args else kwargs.get("url", "")
|
|
1464
|
+
if "/info" in url or "/ping" in url:
|
|
1465
|
+
mock_resp = AsyncMock()
|
|
1466
|
+
mock_resp.status = 200
|
|
1467
|
+
mock_resp.json = AsyncMock(
|
|
1468
|
+
return_value={"archive": True}
|
|
1469
|
+
)
|
|
1470
|
+
mock_resp.text = AsyncMock(return_value="")
|
|
1471
|
+
# Configure as async context manager
|
|
1472
|
+
mock_resp.__aenter__ = AsyncMock(
|
|
1473
|
+
return_value=mock_resp
|
|
1474
|
+
)
|
|
1475
|
+
mock_resp.__aexit__ = AsyncMock(return_value=None)
|
|
1476
|
+
return mock_resp
|
|
1477
|
+
else:
|
|
1478
|
+
# For /download endpoint, return awaitable that resolves to response
|
|
1479
|
+
async def get_download_response():
|
|
1480
|
+
mock_resp = AsyncMock()
|
|
1481
|
+
mock_resp.status = 200
|
|
1482
|
+
mock_resp.headers = {
|
|
1483
|
+
"Content-Type": "application/zip",
|
|
1484
|
+
"Content-Disposition": 'attachment; filename="archive.zip"',
|
|
1485
|
+
}
|
|
1486
|
+
|
|
1487
|
+
async def chunk_iter():
|
|
1488
|
+
yield b"zip data"
|
|
1489
|
+
yield b""
|
|
1490
|
+
|
|
1491
|
+
mock_resp.content.iter_chunked = (
|
|
1492
|
+
lambda size: chunk_iter()
|
|
1493
|
+
)
|
|
1494
|
+
mock_resp.close = Mock()
|
|
1495
|
+
return mock_resp
|
|
1496
|
+
|
|
1497
|
+
return get_download_response()
|
|
1498
|
+
|
|
1499
|
+
mock_session_instance.get = Mock(side_effect=mock_get)
|
|
1500
|
+
|
|
1501
|
+
# Mock retry_request to actually call the function passed to it
|
|
1502
|
+
async def mock_retry(func, *args, **kwargs):
|
|
1503
|
+
return await func(*args, **kwargs)
|
|
1504
|
+
|
|
1505
|
+
with patch(
|
|
1506
|
+
"proximl.utils.transfer.retry_request",
|
|
1507
|
+
side_effect=mock_retry,
|
|
1508
|
+
):
|
|
1509
|
+
mock_file_context = AsyncMock()
|
|
1510
|
+
mock_file_context.__aenter__ = AsyncMock(
|
|
1511
|
+
return_value=AsyncMock()
|
|
1512
|
+
)
|
|
1513
|
+
mock_file_context.__aexit__ = AsyncMock(return_value=None)
|
|
1514
|
+
mock_file_context.__aenter__.return_value.write = (
|
|
1515
|
+
AsyncMock()
|
|
1516
|
+
)
|
|
1517
|
+
with patch(
|
|
1518
|
+
"aiofiles.open", return_value=mock_file_context
|
|
1519
|
+
) as mock_file:
|
|
1520
|
+
mock_finalize_response = AsyncMock()
|
|
1521
|
+
mock_finalize_response.status = 200
|
|
1522
|
+
mock_finalize_response.json = AsyncMock(
|
|
1523
|
+
return_value={"status": "ok"}
|
|
1524
|
+
)
|
|
1525
|
+
mock_finalize_response.__aenter__ = AsyncMock(
|
|
1526
|
+
return_value=mock_finalize_response
|
|
1527
|
+
)
|
|
1528
|
+
mock_finalize_response.__aexit__ = AsyncMock(
|
|
1529
|
+
return_value=None
|
|
1530
|
+
)
|
|
1531
|
+
|
|
1532
|
+
# session.post() should return something that is both awaitable and an async context manager
|
|
1533
|
+
class AwaitableContextManager:
|
|
1534
|
+
def __init__(self, return_value):
|
|
1535
|
+
self.return_value = return_value
|
|
1536
|
+
|
|
1537
|
+
def __await__(self):
|
|
1538
|
+
yield
|
|
1539
|
+
return self
|
|
1540
|
+
|
|
1541
|
+
async def __aenter__(self):
|
|
1542
|
+
return self.return_value
|
|
1543
|
+
|
|
1544
|
+
async def __aexit__(self, *args):
|
|
1545
|
+
return None
|
|
1546
|
+
|
|
1547
|
+
mock_post_context = AwaitableContextManager(
|
|
1548
|
+
mock_finalize_response
|
|
1549
|
+
)
|
|
1550
|
+
mock_session_instance.post = Mock(
|
|
1551
|
+
return_value=mock_post_context
|
|
1552
|
+
)
|
|
1553
|
+
|
|
1554
|
+
with patch(
|
|
1555
|
+
"proximl.utils.transfer.ping_endpoint",
|
|
1556
|
+
new_callable=AsyncMock,
|
|
1557
|
+
):
|
|
1558
|
+
await specimen.download(
|
|
1559
|
+
"example.com", "token", tmpdir, "test.zip"
|
|
1560
|
+
)
|
|
1561
|
+
|
|
1562
|
+
@mark.asyncio
|
|
1563
|
+
async def test_download_info_endpoint_404_fallback(self):
|
|
1564
|
+
with tempfile.TemporaryDirectory() as tmpdir:
|
|
1565
|
+
with patch("aiohttp.ClientSession") as mock_session:
|
|
1566
|
+
mock_session_instance = AsyncMock()
|
|
1567
|
+
mock_session.return_value.__aenter__ = AsyncMock(
|
|
1568
|
+
return_value=mock_session_instance
|
|
1569
|
+
)
|
|
1570
|
+
mock_session.return_value.__aexit__ = AsyncMock(
|
|
1571
|
+
return_value=None
|
|
1572
|
+
)
|
|
1573
|
+
|
|
1574
|
+
# Mock /info endpoint returning 404
|
|
1575
|
+
def mock_get(*args, **kwargs):
|
|
1576
|
+
url = args[0] if args else kwargs.get("url", "")
|
|
1577
|
+
# Handle /ping endpoint for ping_endpoint
|
|
1578
|
+
if "/ping" in url:
|
|
1579
|
+
mock_resp = AsyncMock()
|
|
1580
|
+
mock_resp.status = 200
|
|
1581
|
+
mock_resp.__aenter__ = AsyncMock(
|
|
1582
|
+
return_value=mock_resp
|
|
1583
|
+
)
|
|
1584
|
+
mock_resp.__aexit__ = AsyncMock(return_value=None)
|
|
1585
|
+
return mock_resp
|
|
1586
|
+
# Handle /info endpoint - returns 404
|
|
1587
|
+
elif "/info" in url:
|
|
1588
|
+
raise ClientResponseError(
|
|
1589
|
+
request_info=Mock(),
|
|
1590
|
+
history=(),
|
|
1591
|
+
status=404,
|
|
1592
|
+
message="Not found",
|
|
1593
|
+
)
|
|
1594
|
+
else:
|
|
1595
|
+
# For /download endpoint, return awaitable that resolves to response
|
|
1596
|
+
async def get_download_response():
|
|
1597
|
+
mock_resp = AsyncMock()
|
|
1598
|
+
mock_resp.status = 200
|
|
1599
|
+
mock_resp.headers = {
|
|
1600
|
+
"Content-Type": "application/x-tar"
|
|
1601
|
+
}
|
|
1602
|
+
|
|
1603
|
+
async def chunk_iter():
|
|
1604
|
+
yield b"tar data"
|
|
1605
|
+
yield b""
|
|
1606
|
+
|
|
1607
|
+
mock_resp.content.iter_chunked = (
|
|
1608
|
+
lambda size: chunk_iter()
|
|
1609
|
+
)
|
|
1610
|
+
mock_resp.close = Mock()
|
|
1611
|
+
return mock_resp
|
|
1612
|
+
|
|
1613
|
+
return get_download_response()
|
|
1614
|
+
|
|
1615
|
+
mock_session_instance.get = Mock(side_effect=mock_get)
|
|
1616
|
+
|
|
1617
|
+
# retry_request is called multiple times: _get_info (raises 404), _download, _finalize
|
|
1618
|
+
call_count = [0]
|
|
1619
|
+
|
|
1620
|
+
async def mock_retry(func, *args, **kwargs):
|
|
1621
|
+
call_count[0] += 1
|
|
1622
|
+
if call_count[0] == 1:
|
|
1623
|
+
# First call is _get_info, which should raise 404
|
|
1624
|
+
raise ClientResponseError(
|
|
1625
|
+
request_info=Mock(),
|
|
1626
|
+
history=(),
|
|
1627
|
+
status=404,
|
|
1628
|
+
message="Not found",
|
|
1629
|
+
)
|
|
1630
|
+
else:
|
|
1631
|
+
# Subsequent calls should proceed normally
|
|
1632
|
+
return await func(*args, **kwargs)
|
|
1633
|
+
|
|
1634
|
+
with patch(
|
|
1635
|
+
"proximl.utils.transfer.retry_request",
|
|
1636
|
+
side_effect=mock_retry,
|
|
1637
|
+
):
|
|
1638
|
+
with patch(
|
|
1639
|
+
"asyncio.create_subprocess_exec"
|
|
1640
|
+
) as mock_subprocess:
|
|
1641
|
+
mock_process = AsyncMock()
|
|
1642
|
+
mock_process.stdin = Mock()
|
|
1643
|
+
mock_process.stdin.write = Mock()
|
|
1644
|
+
mock_process.stdin.drain = AsyncMock()
|
|
1645
|
+
mock_process.stdin.close = Mock()
|
|
1646
|
+
mock_process.returncode = 0
|
|
1647
|
+
mock_process.wait = AsyncMock(return_value=0)
|
|
1648
|
+
mock_process.stderr.read = AsyncMock(return_value=b"")
|
|
1649
|
+
mock_subprocess.return_value = mock_process
|
|
1650
|
+
|
|
1651
|
+
mock_finalize_response = AsyncMock()
|
|
1652
|
+
mock_finalize_response.status = 200
|
|
1653
|
+
mock_finalize_response.json = AsyncMock(
|
|
1654
|
+
return_value={"status": "ok"}
|
|
1655
|
+
)
|
|
1656
|
+
mock_finalize_response.__aenter__ = AsyncMock(
|
|
1657
|
+
return_value=mock_finalize_response
|
|
1658
|
+
)
|
|
1659
|
+
mock_finalize_response.__aexit__ = AsyncMock(
|
|
1660
|
+
return_value=None
|
|
1661
|
+
)
|
|
1662
|
+
|
|
1663
|
+
# session.post() should return something that is both awaitable and an async context manager
|
|
1664
|
+
class AwaitableContextManager:
|
|
1665
|
+
def __init__(self, return_value):
|
|
1666
|
+
self.return_value = return_value
|
|
1667
|
+
|
|
1668
|
+
def __await__(self):
|
|
1669
|
+
yield
|
|
1670
|
+
return self
|
|
1671
|
+
|
|
1672
|
+
async def __aenter__(self):
|
|
1673
|
+
return self.return_value
|
|
1674
|
+
|
|
1675
|
+
async def __aexit__(self, *args):
|
|
1676
|
+
return None
|
|
1677
|
+
|
|
1678
|
+
mock_post_context = AwaitableContextManager(
|
|
1679
|
+
mock_finalize_response
|
|
1680
|
+
)
|
|
1681
|
+
mock_session_instance.post = Mock(
|
|
1682
|
+
return_value=mock_post_context
|
|
1683
|
+
)
|
|
1684
|
+
|
|
1685
|
+
with patch(
|
|
1686
|
+
"proximl.utils.transfer.ping_endpoint",
|
|
1687
|
+
new_callable=AsyncMock,
|
|
1688
|
+
):
|
|
1689
|
+
await specimen.download(
|
|
1690
|
+
"example.com", "token", tmpdir
|
|
1691
|
+
)
|
|
1692
|
+
|
|
1693
|
+
@mark.asyncio
|
|
1694
|
+
async def test_download_info_endpoint_connection_error_404(self):
|
|
1695
|
+
with tempfile.TemporaryDirectory() as tmpdir:
|
|
1696
|
+
with patch("aiohttp.ClientSession") as mock_session:
|
|
1697
|
+
mock_session_instance = AsyncMock()
|
|
1698
|
+
mock_session.return_value.__aenter__ = AsyncMock(
|
|
1699
|
+
return_value=mock_session_instance
|
|
1700
|
+
)
|
|
1701
|
+
mock_session.return_value.__aexit__ = AsyncMock(
|
|
1702
|
+
return_value=None
|
|
1703
|
+
)
|
|
1704
|
+
|
|
1705
|
+
# Mock /info endpoint returning ConnectionError with 404
|
|
1706
|
+
def mock_get(*args, **kwargs):
|
|
1707
|
+
url = args[0] if args else kwargs.get("url", "")
|
|
1708
|
+
# Handle /ping endpoint for ping_endpoint
|
|
1709
|
+
if "/ping" in url:
|
|
1710
|
+
mock_resp = AsyncMock()
|
|
1711
|
+
mock_resp.status = 200
|
|
1712
|
+
mock_resp.__aenter__ = AsyncMock(
|
|
1713
|
+
return_value=mock_resp
|
|
1714
|
+
)
|
|
1715
|
+
mock_resp.__aexit__ = AsyncMock(return_value=None)
|
|
1716
|
+
return mock_resp
|
|
1717
|
+
# Handle /info endpoint - raises ConnectionError
|
|
1718
|
+
elif "/info" in url:
|
|
1719
|
+
raise ConnectionError(
|
|
1720
|
+
"Failed to get server info (status 404): Not found"
|
|
1721
|
+
)
|
|
1722
|
+
else:
|
|
1723
|
+
# For /download endpoint, return awaitable that resolves to response
|
|
1724
|
+
async def get_download_response():
|
|
1725
|
+
mock_resp = AsyncMock()
|
|
1726
|
+
mock_resp.status = 200
|
|
1727
|
+
mock_resp.headers = {
|
|
1728
|
+
"Content-Type": "application/x-tar"
|
|
1729
|
+
}
|
|
1730
|
+
|
|
1731
|
+
async def chunk_iter():
|
|
1732
|
+
yield b"tar data"
|
|
1733
|
+
yield b""
|
|
1734
|
+
|
|
1735
|
+
mock_resp.content.iter_chunked = (
|
|
1736
|
+
lambda size: chunk_iter()
|
|
1737
|
+
)
|
|
1738
|
+
mock_resp.close = Mock()
|
|
1739
|
+
return mock_resp
|
|
1740
|
+
|
|
1741
|
+
return get_download_response()
|
|
1742
|
+
|
|
1743
|
+
mock_session_instance.get = Mock(side_effect=mock_get)
|
|
1744
|
+
|
|
1745
|
+
# retry_request is called multiple times: _get_info (raises ConnectionError), _download, _finalize
|
|
1746
|
+
call_count = [0]
|
|
1747
|
+
|
|
1748
|
+
async def mock_retry(func, *args, **kwargs):
|
|
1749
|
+
call_count[0] += 1
|
|
1750
|
+
if call_count[0] == 1:
|
|
1751
|
+
# First call is _get_info, which should raise ConnectionError
|
|
1752
|
+
raise ConnectionError(
|
|
1753
|
+
"Failed to get server info (status 404): Not found"
|
|
1754
|
+
)
|
|
1755
|
+
else:
|
|
1756
|
+
# Subsequent calls should proceed normally
|
|
1757
|
+
return await func(*args, **kwargs)
|
|
1758
|
+
|
|
1759
|
+
with patch(
|
|
1760
|
+
"proximl.utils.transfer.retry_request",
|
|
1761
|
+
side_effect=mock_retry,
|
|
1762
|
+
):
|
|
1763
|
+
with patch(
|
|
1764
|
+
"asyncio.create_subprocess_exec"
|
|
1765
|
+
) as mock_subprocess:
|
|
1766
|
+
mock_process = AsyncMock()
|
|
1767
|
+
mock_process.stdin = Mock()
|
|
1768
|
+
mock_process.stdin.write = Mock()
|
|
1769
|
+
mock_process.stdin.drain = AsyncMock()
|
|
1770
|
+
mock_process.stdin.close = Mock()
|
|
1771
|
+
mock_process.returncode = 0
|
|
1772
|
+
mock_process.wait = AsyncMock(return_value=0)
|
|
1773
|
+
# Return empty stderr to avoid tar error message
|
|
1774
|
+
mock_process.stderr.read = AsyncMock(return_value=b"")
|
|
1775
|
+
mock_subprocess.return_value = mock_process
|
|
1776
|
+
# Ensure stdin is properly set up
|
|
1777
|
+
if (
|
|
1778
|
+
not hasattr(mock_process, "stdin")
|
|
1779
|
+
or mock_process.stdin is None
|
|
1780
|
+
):
|
|
1781
|
+
mock_process.stdin = Mock()
|
|
1782
|
+
mock_process.stdin.write = Mock()
|
|
1783
|
+
mock_process.stdin.drain = AsyncMock()
|
|
1784
|
+
mock_process.stdin.close = Mock()
|
|
1785
|
+
|
|
1786
|
+
mock_finalize_response = AsyncMock()
|
|
1787
|
+
mock_finalize_response.status = 200
|
|
1788
|
+
mock_finalize_response.json = AsyncMock(
|
|
1789
|
+
return_value={"status": "ok"}
|
|
1790
|
+
)
|
|
1791
|
+
|
|
1792
|
+
mock_finalize_response.__aenter__ = AsyncMock(
|
|
1793
|
+
return_value=mock_finalize_response
|
|
1794
|
+
)
|
|
1795
|
+
|
|
1796
|
+
mock_finalize_response.__aexit__ = AsyncMock(
|
|
1797
|
+
return_value=None
|
|
1798
|
+
)
|
|
1799
|
+
|
|
1800
|
+
# session.post() should return something that is both awaitable and an async context manager
|
|
1801
|
+
class AwaitableContextManager:
|
|
1802
|
+
def __init__(self, return_value):
|
|
1803
|
+
self.return_value = return_value
|
|
1804
|
+
|
|
1805
|
+
def __await__(self):
|
|
1806
|
+
yield
|
|
1807
|
+
return self
|
|
1808
|
+
|
|
1809
|
+
async def __aenter__(self):
|
|
1810
|
+
return self.return_value
|
|
1811
|
+
|
|
1812
|
+
async def __aexit__(self, *args):
|
|
1813
|
+
return None
|
|
1814
|
+
|
|
1815
|
+
mock_post_context = AwaitableContextManager(
|
|
1816
|
+
mock_finalize_response
|
|
1817
|
+
)
|
|
1818
|
+
mock_session_instance.post = Mock(
|
|
1819
|
+
return_value=mock_post_context
|
|
1820
|
+
)
|
|
1821
|
+
|
|
1822
|
+
with patch(
|
|
1823
|
+
"proximl.utils.transfer.ping_endpoint",
|
|
1824
|
+
new_callable=AsyncMock,
|
|
1825
|
+
):
|
|
1826
|
+
await specimen.download(
|
|
1827
|
+
"example.com", "token", tmpdir
|
|
1828
|
+
)
|
|
1829
|
+
|
|
1830
|
+
@mark.asyncio
|
|
1831
|
+
async def test_download_info_endpoint_non_404_error(self):
|
|
1832
|
+
with tempfile.TemporaryDirectory() as tmpdir:
|
|
1833
|
+
with patch("aiohttp.ClientSession") as mock_session:
|
|
1834
|
+
mock_session_instance = AsyncMock()
|
|
1835
|
+
mock_session.return_value.__aenter__ = AsyncMock(
|
|
1836
|
+
return_value=mock_session_instance
|
|
1837
|
+
)
|
|
1838
|
+
mock_session.return_value.__aexit__ = AsyncMock(
|
|
1839
|
+
return_value=None
|
|
1840
|
+
)
|
|
1841
|
+
|
|
1842
|
+
# Mock /info endpoint returning non-404 error
|
|
1843
|
+
with patch(
|
|
1844
|
+
"proximl.utils.transfer.retry_request",
|
|
1845
|
+
side_effect=ConnectionError(
|
|
1846
|
+
"Failed to get server info (status 500): Server error"
|
|
1847
|
+
),
|
|
1848
|
+
):
|
|
1849
|
+
with patch(
|
|
1850
|
+
"proximl.utils.transfer.ping_endpoint",
|
|
1851
|
+
new_callable=AsyncMock,
|
|
1852
|
+
):
|
|
1853
|
+
with raises(ConnectionError):
|
|
1854
|
+
await specimen.download(
|
|
1855
|
+
"example.com", "token", tmpdir
|
|
1856
|
+
)
|
|
1857
|
+
|
|
1858
|
+
@mark.asyncio
|
|
1859
|
+
async def test_download_info_endpoint_invalid_url(self):
|
|
1860
|
+
with tempfile.TemporaryDirectory() as tmpdir:
|
|
1861
|
+
with patch("aiohttp.ClientSession") as mock_session:
|
|
1862
|
+
mock_session_instance = AsyncMock()
|
|
1863
|
+
mock_session.return_value.__aenter__ = AsyncMock(
|
|
1864
|
+
return_value=mock_session_instance
|
|
1865
|
+
)
|
|
1866
|
+
mock_session.return_value.__aexit__ = AsyncMock(
|
|
1867
|
+
return_value=None
|
|
1868
|
+
)
|
|
1869
|
+
|
|
1870
|
+
# Mock /info endpoint returning InvalidURL
|
|
1871
|
+
with patch(
|
|
1872
|
+
"proximl.utils.transfer.retry_request",
|
|
1873
|
+
side_effect=InvalidURL("Invalid URL"),
|
|
1874
|
+
):
|
|
1875
|
+
with patch(
|
|
1876
|
+
"proximl.utils.transfer.ping_endpoint",
|
|
1877
|
+
new_callable=AsyncMock,
|
|
1878
|
+
):
|
|
1879
|
+
with raises(
|
|
1880
|
+
ConnectionError, match="Invalid endpoint URL"
|
|
1881
|
+
):
|
|
1882
|
+
await specimen.download(
|
|
1883
|
+
"example.com", "token", tmpdir
|
|
1884
|
+
)
|
|
1885
|
+
|
|
1886
|
+
@mark.asyncio
|
|
1887
|
+
async def test_download_info_endpoint_error_reading_body(self):
|
|
1888
|
+
with tempfile.TemporaryDirectory() as tmpdir:
|
|
1889
|
+
with patch("aiohttp.ClientSession") as mock_session:
|
|
1890
|
+
mock_session_instance = AsyncMock()
|
|
1891
|
+
mock_session.return_value.__aenter__ = AsyncMock(
|
|
1892
|
+
return_value=mock_session_instance
|
|
1893
|
+
)
|
|
1894
|
+
mock_session.return_value.__aexit__ = AsyncMock(
|
|
1895
|
+
return_value=None
|
|
1896
|
+
)
|
|
1897
|
+
|
|
1898
|
+
# Mock /info endpoint with error reading response body
|
|
1899
|
+
# This tests the exception handling in _get_info when response.text() raises
|
|
1900
|
+
async def _get_info_with_error(*args, **kwargs):
|
|
1901
|
+
mock_resp = AsyncMock()
|
|
1902
|
+
mock_resp.status = 500
|
|
1903
|
+
mock_resp.text = AsyncMock(
|
|
1904
|
+
side_effect=Exception("Read error")
|
|
1905
|
+
)
|
|
1906
|
+
mock_resp.request_info = Mock()
|
|
1907
|
+
mock_resp.history = ()
|
|
1908
|
+
mock_resp.__aenter__ = AsyncMock(return_value=mock_resp)
|
|
1909
|
+
mock_resp.__aexit__ = AsyncMock(return_value=None)
|
|
1910
|
+
async with mock_resp:
|
|
1911
|
+
if mock_resp.status != 200:
|
|
1912
|
+
try:
|
|
1913
|
+
error_text = await mock_resp.text()
|
|
1914
|
+
except Exception:
|
|
1915
|
+
error_text = f"Unable to read response body (status: {mock_resp.status})"
|
|
1916
|
+
raise ConnectionError(
|
|
1917
|
+
f"Failed to get server info (status {mock_resp.status}): {error_text}"
|
|
1918
|
+
)
|
|
1919
|
+
|
|
1920
|
+
with patch(
|
|
1921
|
+
"proximl.utils.transfer.retry_request",
|
|
1922
|
+
side_effect=_get_info_with_error,
|
|
1923
|
+
):
|
|
1924
|
+
with patch(
|
|
1925
|
+
"proximl.utils.transfer.ping_endpoint",
|
|
1926
|
+
new_callable=AsyncMock,
|
|
1927
|
+
):
|
|
1928
|
+
with raises(
|
|
1929
|
+
ConnectionError, match="Failed to get server info"
|
|
1930
|
+
):
|
|
1931
|
+
await specimen.download(
|
|
1932
|
+
"example.com", "token", tmpdir
|
|
1933
|
+
)
|
|
1934
|
+
|
|
1935
|
+
@mark.asyncio
|
|
1936
|
+
async def test_download_zip_content_type_fallback(self):
|
|
1937
|
+
with tempfile.TemporaryDirectory() as tmpdir:
|
|
1938
|
+
with patch("aiohttp.ClientSession") as mock_session:
|
|
1939
|
+
mock_session_instance = AsyncMock()
|
|
1940
|
+
mock_session.return_value.__aenter__ = AsyncMock(
|
|
1941
|
+
return_value=mock_session_instance
|
|
1942
|
+
)
|
|
1943
|
+
mock_session.return_value.__aexit__ = AsyncMock(
|
|
1944
|
+
return_value=None
|
|
1945
|
+
)
|
|
1946
|
+
|
|
1947
|
+
# Mock /info endpoint returning TAR mode but Content-Type says zip
|
|
1948
|
+
def mock_get(*args, **kwargs):
|
|
1949
|
+
url = args[0] if args else kwargs.get("url", "")
|
|
1950
|
+
if "/info" in url or "/ping" in url:
|
|
1951
|
+
mock_resp = AsyncMock()
|
|
1952
|
+
mock_resp.status = 200
|
|
1953
|
+
mock_resp.json = AsyncMock(
|
|
1954
|
+
return_value={"archive": False}
|
|
1955
|
+
)
|
|
1956
|
+
mock_resp.text = AsyncMock(return_value="")
|
|
1957
|
+
mock_resp.__aenter__ = AsyncMock(
|
|
1958
|
+
return_value=mock_resp
|
|
1959
|
+
)
|
|
1960
|
+
mock_resp.__aexit__ = AsyncMock(return_value=None)
|
|
1961
|
+
return mock_resp
|
|
1962
|
+
else:
|
|
1963
|
+
# For /download endpoint, return awaitable that resolves to response
|
|
1964
|
+
async def get_download_response():
|
|
1965
|
+
mock_resp = AsyncMock()
|
|
1966
|
+
mock_resp.status = 200
|
|
1967
|
+
mock_resp.headers = {
|
|
1968
|
+
"Content-Type": "application/zip",
|
|
1969
|
+
"Content-Disposition": 'attachment; filename="archive.zip"',
|
|
1970
|
+
}
|
|
1971
|
+
|
|
1972
|
+
async def chunk_iter():
|
|
1973
|
+
yield b"zip data"
|
|
1974
|
+
yield b""
|
|
1975
|
+
|
|
1976
|
+
mock_resp.content.iter_chunked = (
|
|
1977
|
+
lambda size: chunk_iter()
|
|
1978
|
+
)
|
|
1979
|
+
mock_resp.close = Mock()
|
|
1980
|
+
return mock_resp
|
|
1981
|
+
|
|
1982
|
+
return get_download_response()
|
|
1983
|
+
|
|
1984
|
+
mock_session_instance.get = Mock(side_effect=mock_get)
|
|
1985
|
+
|
|
1986
|
+
# retry_request is called 3 times: _get_info, _download, _finalize
|
|
1987
|
+
async def mock_retry(func, *args, **kwargs):
|
|
1988
|
+
return await func(*args, **kwargs)
|
|
1989
|
+
|
|
1990
|
+
with patch(
|
|
1991
|
+
"proximl.utils.transfer.retry_request",
|
|
1992
|
+
side_effect=mock_retry,
|
|
1993
|
+
):
|
|
1994
|
+
mock_file_context = AsyncMock()
|
|
1995
|
+
mock_file_context.__aenter__ = AsyncMock(
|
|
1996
|
+
return_value=AsyncMock()
|
|
1997
|
+
)
|
|
1998
|
+
mock_file_context.__aexit__ = AsyncMock(return_value=None)
|
|
1999
|
+
mock_file_context.__aenter__.return_value.write = (
|
|
2000
|
+
AsyncMock()
|
|
2001
|
+
)
|
|
2002
|
+
with patch(
|
|
2003
|
+
"aiofiles.open", return_value=mock_file_context
|
|
2004
|
+
) as mock_file:
|
|
2005
|
+
mock_finalize_response = AsyncMock()
|
|
2006
|
+
mock_finalize_response.status = 200
|
|
2007
|
+
mock_finalize_response.json = AsyncMock(
|
|
2008
|
+
return_value={"status": "ok"}
|
|
2009
|
+
)
|
|
2010
|
+
|
|
2011
|
+
mock_finalize_response.__aenter__ = AsyncMock(
|
|
2012
|
+
return_value=mock_finalize_response
|
|
2013
|
+
)
|
|
2014
|
+
|
|
2015
|
+
mock_finalize_response.__aexit__ = AsyncMock(
|
|
2016
|
+
return_value=None
|
|
2017
|
+
)
|
|
2018
|
+
|
|
2019
|
+
# session.post() should return something that is both awaitable and an async context manager
|
|
2020
|
+
class AwaitableContextManager:
|
|
2021
|
+
def __init__(self, return_value):
|
|
2022
|
+
self.return_value = return_value
|
|
2023
|
+
|
|
2024
|
+
def __await__(self):
|
|
2025
|
+
yield
|
|
2026
|
+
return self
|
|
2027
|
+
|
|
2028
|
+
async def __aenter__(self):
|
|
2029
|
+
return self.return_value
|
|
2030
|
+
|
|
2031
|
+
async def __aexit__(self, *args):
|
|
2032
|
+
return None
|
|
2033
|
+
|
|
2034
|
+
mock_post_context = AwaitableContextManager(
|
|
2035
|
+
mock_finalize_response
|
|
2036
|
+
)
|
|
2037
|
+
mock_session_instance.post = Mock(
|
|
2038
|
+
return_value=mock_post_context
|
|
2039
|
+
)
|
|
2040
|
+
|
|
2041
|
+
with patch(
|
|
2042
|
+
"proximl.utils.transfer.ping_endpoint",
|
|
2043
|
+
new_callable=AsyncMock,
|
|
2044
|
+
):
|
|
2045
|
+
await specimen.download(
|
|
2046
|
+
"example.com", "token", tmpdir
|
|
2047
|
+
)
|
|
2048
|
+
|
|
2049
|
+
@mark.asyncio
|
|
2050
|
+
async def test_download_content_length_logging(self):
|
|
2051
|
+
with tempfile.TemporaryDirectory() as tmpdir:
|
|
2052
|
+
with patch("aiohttp.ClientSession") as mock_session:
|
|
2053
|
+
mock_session_instance = AsyncMock()
|
|
2054
|
+
mock_session.return_value.__aenter__ = AsyncMock(
|
|
2055
|
+
return_value=mock_session_instance
|
|
2056
|
+
)
|
|
2057
|
+
mock_session.return_value.__aexit__ = AsyncMock(
|
|
2058
|
+
return_value=None
|
|
2059
|
+
)
|
|
2060
|
+
|
|
2061
|
+
def mock_get(*args, **kwargs):
|
|
2062
|
+
url = args[0] if args else kwargs.get("url", "")
|
|
2063
|
+
if "/info" in url or "/ping" in url:
|
|
2064
|
+
mock_resp = AsyncMock()
|
|
2065
|
+
mock_resp.status = 200
|
|
2066
|
+
mock_resp.json = AsyncMock(
|
|
2067
|
+
return_value={"archive": False}
|
|
2068
|
+
)
|
|
2069
|
+
mock_resp.text = AsyncMock(return_value="")
|
|
2070
|
+
mock_resp.__aenter__ = AsyncMock(
|
|
2071
|
+
return_value=mock_resp
|
|
2072
|
+
)
|
|
2073
|
+
mock_resp.__aexit__ = AsyncMock(return_value=None)
|
|
2074
|
+
return mock_resp
|
|
2075
|
+
else:
|
|
2076
|
+
# For /download endpoint, return awaitable that resolves to response
|
|
2077
|
+
async def get_download_response():
|
|
2078
|
+
mock_resp = AsyncMock()
|
|
2079
|
+
mock_resp.status = 200
|
|
2080
|
+
mock_resp.headers = {
|
|
2081
|
+
"Content-Type": "application/x-tar",
|
|
2082
|
+
"Content-Length": "1024",
|
|
2083
|
+
}
|
|
2084
|
+
|
|
2085
|
+
async def chunk_iter():
|
|
2086
|
+
yield b"tar data"
|
|
2087
|
+
yield b""
|
|
2088
|
+
|
|
2089
|
+
mock_resp.content.iter_chunked = (
|
|
2090
|
+
lambda size: chunk_iter()
|
|
2091
|
+
)
|
|
2092
|
+
mock_resp.close = Mock()
|
|
2093
|
+
return mock_resp
|
|
2094
|
+
|
|
2095
|
+
return get_download_response()
|
|
2096
|
+
|
|
2097
|
+
mock_session_instance.get = Mock(side_effect=mock_get)
|
|
2098
|
+
|
|
2099
|
+
# retry_request is called 3 times: _get_info, _download, _finalize
|
|
2100
|
+
async def mock_retry(func, *args, **kwargs):
|
|
2101
|
+
return await func(*args, **kwargs)
|
|
2102
|
+
|
|
2103
|
+
with patch(
|
|
2104
|
+
"proximl.utils.transfer.retry_request",
|
|
2105
|
+
side_effect=mock_retry,
|
|
2106
|
+
):
|
|
2107
|
+
with patch(
|
|
2108
|
+
"asyncio.create_subprocess_exec"
|
|
2109
|
+
) as mock_subprocess:
|
|
2110
|
+
mock_process = AsyncMock()
|
|
2111
|
+
mock_process.stdin = Mock()
|
|
2112
|
+
mock_process.stdin.write = Mock()
|
|
2113
|
+
mock_process.stdin.drain = AsyncMock()
|
|
2114
|
+
mock_process.stdin.close = Mock()
|
|
2115
|
+
mock_process.returncode = 0
|
|
2116
|
+
mock_process.wait = AsyncMock(return_value=0)
|
|
2117
|
+
mock_process.stderr.read = AsyncMock(return_value=b"")
|
|
2118
|
+
mock_subprocess.return_value = mock_process
|
|
2119
|
+
|
|
2120
|
+
mock_finalize_response = AsyncMock()
|
|
2121
|
+
mock_finalize_response.status = 200
|
|
2122
|
+
mock_finalize_response.json = AsyncMock(
|
|
2123
|
+
return_value={"status": "ok"}
|
|
2124
|
+
)
|
|
2125
|
+
mock_finalize_response.__aenter__ = AsyncMock(
|
|
2126
|
+
return_value=mock_finalize_response
|
|
2127
|
+
)
|
|
2128
|
+
mock_finalize_response.__aexit__ = AsyncMock(
|
|
2129
|
+
return_value=None
|
|
2130
|
+
)
|
|
2131
|
+
|
|
2132
|
+
# session.post() should return something that is both awaitable and an async context manager
|
|
2133
|
+
class AwaitableContextManager:
|
|
2134
|
+
def __init__(self, return_value):
|
|
2135
|
+
self.return_value = return_value
|
|
2136
|
+
|
|
2137
|
+
def __await__(self):
|
|
2138
|
+
yield
|
|
2139
|
+
return self
|
|
2140
|
+
|
|
2141
|
+
async def __aenter__(self):
|
|
2142
|
+
return self.return_value
|
|
2143
|
+
|
|
2144
|
+
async def __aexit__(self, *args):
|
|
2145
|
+
return None
|
|
2146
|
+
|
|
2147
|
+
mock_post_context = AwaitableContextManager(
|
|
2148
|
+
mock_finalize_response
|
|
2149
|
+
)
|
|
2150
|
+
mock_session_instance.post = Mock(
|
|
2151
|
+
return_value=mock_post_context
|
|
2152
|
+
)
|
|
2153
|
+
|
|
2154
|
+
with patch(
|
|
2155
|
+
"proximl.utils.transfer.ping_endpoint",
|
|
2156
|
+
new_callable=AsyncMock,
|
|
2157
|
+
):
|
|
2158
|
+
await specimen.download(
|
|
2159
|
+
"example.com", "token", tmpdir
|
|
2160
|
+
)
|
|
2161
|
+
|
|
2162
|
+
@mark.asyncio
|
|
2163
|
+
async def test_download_empty_file_error(self):
|
|
2164
|
+
with tempfile.TemporaryDirectory() as tmpdir:
|
|
2165
|
+
with patch("aiohttp.ClientSession") as mock_session:
|
|
2166
|
+
mock_session_instance = AsyncMock()
|
|
2167
|
+
mock_session.return_value.__aenter__ = AsyncMock(
|
|
2168
|
+
return_value=mock_session_instance
|
|
2169
|
+
)
|
|
2170
|
+
mock_session.return_value.__aexit__ = AsyncMock(
|
|
2171
|
+
return_value=None
|
|
2172
|
+
)
|
|
2173
|
+
|
|
2174
|
+
def mock_get(*args, **kwargs):
|
|
2175
|
+
url = args[0] if args else kwargs.get("url", "")
|
|
2176
|
+
if "/info" in url or "/ping" in url:
|
|
2177
|
+
mock_resp = AsyncMock()
|
|
2178
|
+
mock_resp.status = 200
|
|
2179
|
+
mock_resp.json = AsyncMock(
|
|
2180
|
+
return_value={"archive": True}
|
|
2181
|
+
)
|
|
2182
|
+
mock_resp.text = AsyncMock(return_value="")
|
|
2183
|
+
mock_resp.__aenter__ = AsyncMock(
|
|
2184
|
+
return_value=mock_resp
|
|
2185
|
+
)
|
|
2186
|
+
mock_resp.__aexit__ = AsyncMock(return_value=None)
|
|
2187
|
+
return mock_resp
|
|
2188
|
+
else:
|
|
2189
|
+
# For /download endpoint, return awaitable that resolves to response
|
|
2190
|
+
async def get_download_response():
|
|
2191
|
+
mock_resp = AsyncMock()
|
|
2192
|
+
mock_resp.status = 200
|
|
2193
|
+
mock_resp.headers = {
|
|
2194
|
+
"Content-Type": "application/zip"
|
|
2195
|
+
}
|
|
2196
|
+
|
|
2197
|
+
async def chunk_iter():
|
|
2198
|
+
return
|
|
2199
|
+
yield # Make it an async generator
|
|
2200
|
+
|
|
2201
|
+
mock_resp.content.iter_chunked = (
|
|
2202
|
+
lambda size: chunk_iter()
|
|
2203
|
+
)
|
|
2204
|
+
mock_resp.close = Mock()
|
|
2205
|
+
return mock_resp
|
|
2206
|
+
|
|
2207
|
+
return get_download_response()
|
|
2208
|
+
|
|
2209
|
+
mock_session_instance.get = Mock(side_effect=mock_get)
|
|
2210
|
+
|
|
2211
|
+
# Mock retry_request to actually call the function passed to it
|
|
2212
|
+
async def mock_retry(func, *args, **kwargs):
|
|
2213
|
+
return await func(*args, **kwargs)
|
|
2214
|
+
|
|
2215
|
+
with patch(
|
|
2216
|
+
"proximl.utils.transfer.retry_request",
|
|
2217
|
+
side_effect=mock_retry,
|
|
2218
|
+
):
|
|
2219
|
+
mock_file_context = AsyncMock()
|
|
2220
|
+
mock_file_context.__aenter__ = AsyncMock(
|
|
2221
|
+
return_value=AsyncMock()
|
|
2222
|
+
)
|
|
2223
|
+
mock_file_context.__aexit__ = AsyncMock(return_value=None)
|
|
2224
|
+
mock_file_context.__aenter__.return_value.write = (
|
|
2225
|
+
AsyncMock()
|
|
2226
|
+
)
|
|
2227
|
+
with patch(
|
|
2228
|
+
"aiofiles.open", return_value=mock_file_context
|
|
2229
|
+
) as mock_file:
|
|
2230
|
+
with patch(
|
|
2231
|
+
"proximl.utils.transfer.ping_endpoint",
|
|
2232
|
+
new_callable=AsyncMock,
|
|
2233
|
+
):
|
|
2234
|
+
with raises(
|
|
2235
|
+
ConnectionError,
|
|
2236
|
+
match="Downloaded file is empty",
|
|
2237
|
+
):
|
|
2238
|
+
await specimen.download(
|
|
2239
|
+
"example.com", "token", tmpdir
|
|
2240
|
+
)
|
|
2241
|
+
|
|
2242
|
+
@mark.asyncio
|
|
2243
|
+
async def test_download_tar_extraction_failure(self):
|
|
2244
|
+
with tempfile.TemporaryDirectory() as tmpdir:
|
|
2245
|
+
with patch("aiohttp.ClientSession") as mock_session:
|
|
2246
|
+
mock_session_instance = AsyncMock()
|
|
2247
|
+
mock_session.return_value.__aenter__ = AsyncMock(
|
|
2248
|
+
return_value=mock_session_instance
|
|
2249
|
+
)
|
|
2250
|
+
mock_session.return_value.__aexit__ = AsyncMock(
|
|
2251
|
+
return_value=None
|
|
2252
|
+
)
|
|
2253
|
+
|
|
2254
|
+
def mock_get(*args, **kwargs):
|
|
2255
|
+
url = args[0] if args else kwargs.get("url", "")
|
|
2256
|
+
if "/info" in url or "/ping" in url:
|
|
2257
|
+
mock_resp = AsyncMock()
|
|
2258
|
+
mock_resp.status = 200
|
|
2259
|
+
mock_resp.json = AsyncMock(
|
|
2260
|
+
return_value={"archive": False}
|
|
2261
|
+
)
|
|
2262
|
+
mock_resp.text = AsyncMock(return_value="")
|
|
2263
|
+
mock_resp.__aenter__ = AsyncMock(
|
|
2264
|
+
return_value=mock_resp
|
|
2265
|
+
)
|
|
2266
|
+
mock_resp.__aexit__ = AsyncMock(return_value=None)
|
|
2267
|
+
return mock_resp
|
|
2268
|
+
else:
|
|
2269
|
+
# For /download endpoint, return awaitable that resolves to response
|
|
2270
|
+
async def get_download_response():
|
|
2271
|
+
mock_resp = AsyncMock()
|
|
2272
|
+
mock_resp.status = 200
|
|
2273
|
+
mock_resp.headers = {
|
|
2274
|
+
"Content-Type": "application/x-tar"
|
|
2275
|
+
}
|
|
2276
|
+
|
|
2277
|
+
async def chunk_iter():
|
|
2278
|
+
yield b"tar data"
|
|
2279
|
+
yield b""
|
|
2280
|
+
|
|
2281
|
+
mock_resp.content.iter_chunked = (
|
|
2282
|
+
lambda size: chunk_iter()
|
|
2283
|
+
)
|
|
2284
|
+
mock_resp.close = Mock()
|
|
2285
|
+
return mock_resp
|
|
2286
|
+
|
|
2287
|
+
return get_download_response()
|
|
2288
|
+
|
|
2289
|
+
mock_session_instance.get = Mock(side_effect=mock_get)
|
|
2290
|
+
|
|
2291
|
+
# retry_request is called 3 times: _get_info, _download, _finalize
|
|
2292
|
+
async def mock_retry(func, *args, **kwargs):
|
|
2293
|
+
return await func(*args, **kwargs)
|
|
2294
|
+
|
|
2295
|
+
with patch(
|
|
2296
|
+
"proximl.utils.transfer.retry_request",
|
|
2297
|
+
side_effect=mock_retry,
|
|
2298
|
+
):
|
|
2299
|
+
with patch(
|
|
2300
|
+
"asyncio.create_subprocess_exec"
|
|
2301
|
+
) as mock_subprocess:
|
|
2302
|
+
mock_process = AsyncMock()
|
|
2303
|
+
mock_process.stdin = Mock()
|
|
2304
|
+
mock_process.stdin.write = Mock()
|
|
2305
|
+
mock_process.stdin.drain = AsyncMock()
|
|
2306
|
+
mock_process.stdin.close = Mock()
|
|
2307
|
+
mock_process.wait = AsyncMock(return_value=1)
|
|
2308
|
+
mock_process.stderr.read = AsyncMock(
|
|
2309
|
+
return_value=b"tar error"
|
|
2310
|
+
)
|
|
2311
|
+
mock_subprocess.return_value = mock_process
|
|
2312
|
+
|
|
2313
|
+
with patch(
|
|
2314
|
+
"proximl.utils.transfer.ping_endpoint",
|
|
2315
|
+
new_callable=AsyncMock,
|
|
2316
|
+
):
|
|
2317
|
+
with raises(
|
|
2318
|
+
ProxiMLException, match="tar extraction failed"
|
|
2319
|
+
):
|
|
2320
|
+
await specimen.download(
|
|
2321
|
+
"example.com", "token", tmpdir
|
|
2322
|
+
)
|
|
2323
|
+
|
|
2324
|
+
@mark.asyncio
|
|
2325
|
+
async def test_download_tar_extraction_failure_no_stderr(self):
|
|
2326
|
+
with tempfile.TemporaryDirectory() as tmpdir:
|
|
2327
|
+
with patch("aiohttp.ClientSession") as mock_session:
|
|
2328
|
+
mock_session_instance = AsyncMock()
|
|
2329
|
+
mock_session.return_value.__aenter__ = AsyncMock(
|
|
2330
|
+
return_value=mock_session_instance
|
|
2331
|
+
)
|
|
2332
|
+
mock_session.return_value.__aexit__ = AsyncMock(
|
|
2333
|
+
return_value=None
|
|
2334
|
+
)
|
|
2335
|
+
|
|
2336
|
+
def mock_get(*args, **kwargs):
|
|
2337
|
+
url = args[0] if args else kwargs.get("url", "")
|
|
2338
|
+
if "/info" in url or "/ping" in url:
|
|
2339
|
+
mock_resp = AsyncMock()
|
|
2340
|
+
mock_resp.status = 200
|
|
2341
|
+
mock_resp.json = AsyncMock(
|
|
2342
|
+
return_value={"archive": False}
|
|
2343
|
+
)
|
|
2344
|
+
mock_resp.text = AsyncMock(return_value="")
|
|
2345
|
+
mock_resp.__aenter__ = AsyncMock(
|
|
2346
|
+
return_value=mock_resp
|
|
2347
|
+
)
|
|
2348
|
+
mock_resp.__aexit__ = AsyncMock(return_value=None)
|
|
2349
|
+
return mock_resp
|
|
2350
|
+
else:
|
|
2351
|
+
# For /download endpoint, return awaitable that resolves to response
|
|
2352
|
+
async def get_download_response():
|
|
2353
|
+
mock_resp = AsyncMock()
|
|
2354
|
+
mock_resp.status = 200
|
|
2355
|
+
mock_resp.headers = {
|
|
2356
|
+
"Content-Type": "application/x-tar"
|
|
2357
|
+
}
|
|
2358
|
+
|
|
2359
|
+
async def chunk_iter():
|
|
2360
|
+
yield b"tar data"
|
|
2361
|
+
yield b""
|
|
2362
|
+
|
|
2363
|
+
mock_resp.content.iter_chunked = (
|
|
2364
|
+
lambda size: chunk_iter()
|
|
2365
|
+
)
|
|
2366
|
+
mock_resp.close = Mock()
|
|
2367
|
+
return mock_resp
|
|
2368
|
+
|
|
2369
|
+
return get_download_response()
|
|
2370
|
+
|
|
2371
|
+
mock_session_instance.get = Mock(side_effect=mock_get)
|
|
2372
|
+
|
|
2373
|
+
# retry_request is called 3 times: _get_info, _download, _finalize
|
|
2374
|
+
async def mock_retry(func, *args, **kwargs):
|
|
2375
|
+
return await func(*args, **kwargs)
|
|
2376
|
+
|
|
2377
|
+
with patch(
|
|
2378
|
+
"proximl.utils.transfer.retry_request",
|
|
2379
|
+
side_effect=mock_retry,
|
|
2380
|
+
):
|
|
2381
|
+
with patch(
|
|
2382
|
+
"asyncio.create_subprocess_exec"
|
|
2383
|
+
) as mock_subprocess:
|
|
2384
|
+
mock_process = AsyncMock()
|
|
2385
|
+
mock_process.stdin = Mock()
|
|
2386
|
+
mock_process.stdin.write = Mock()
|
|
2387
|
+
mock_process.stdin.drain = AsyncMock()
|
|
2388
|
+
mock_process.stdin.close = Mock()
|
|
2389
|
+
mock_process.returncode = 1
|
|
2390
|
+
mock_process.wait = AsyncMock(return_value=1)
|
|
2391
|
+
mock_process.stderr.read = AsyncMock(return_value=None)
|
|
2392
|
+
mock_subprocess.return_value = mock_process
|
|
2393
|
+
|
|
2394
|
+
with patch(
|
|
2395
|
+
"proximl.utils.transfer.ping_endpoint",
|
|
2396
|
+
new_callable=AsyncMock,
|
|
2397
|
+
):
|
|
2398
|
+
with raises(
|
|
2399
|
+
ProxiMLException, match="tar extraction failed"
|
|
2400
|
+
):
|
|
2401
|
+
await specimen.download(
|
|
2402
|
+
"example.com", "token", tmpdir
|
|
2403
|
+
)
|
|
2404
|
+
|
|
2405
|
+
@mark.asyncio
|
|
2406
|
+
async def test_download_404_error(self):
|
|
2407
|
+
with tempfile.TemporaryDirectory() as tmpdir:
|
|
2408
|
+
with patch("aiohttp.ClientSession") as mock_session:
|
|
2409
|
+
mock_session_instance = AsyncMock()
|
|
2410
|
+
mock_session.return_value.__aenter__ = AsyncMock(
|
|
2411
|
+
return_value=mock_session_instance
|
|
2412
|
+
)
|
|
2413
|
+
mock_session.return_value.__aexit__ = AsyncMock(
|
|
2414
|
+
return_value=None
|
|
2415
|
+
)
|
|
2416
|
+
|
|
2417
|
+
def mock_get(*args, **kwargs):
|
|
2418
|
+
url = args[0] if args else kwargs.get("url", "")
|
|
2419
|
+
if "/info" in url or "/ping" in url:
|
|
2420
|
+
mock_resp = AsyncMock()
|
|
2421
|
+
mock_resp.status = 200
|
|
2422
|
+
mock_resp.json = AsyncMock(
|
|
2423
|
+
return_value={"archive": False}
|
|
2424
|
+
)
|
|
2425
|
+
mock_resp.text = AsyncMock(return_value="")
|
|
2426
|
+
mock_resp.__aenter__ = AsyncMock(
|
|
2427
|
+
return_value=mock_resp
|
|
2428
|
+
)
|
|
2429
|
+
mock_resp.__aexit__ = AsyncMock(return_value=None)
|
|
2430
|
+
return mock_resp
|
|
2431
|
+
else:
|
|
2432
|
+
# For /download endpoint, return awaitable that raises 404
|
|
2433
|
+
async def get_download_response():
|
|
2434
|
+
raise ClientResponseError(
|
|
2435
|
+
request_info=Mock(),
|
|
2436
|
+
history=(),
|
|
2437
|
+
status=404,
|
|
2438
|
+
message="Not found",
|
|
2439
|
+
)
|
|
2440
|
+
|
|
2441
|
+
return get_download_response()
|
|
2442
|
+
|
|
2443
|
+
mock_session_instance.get = Mock(side_effect=mock_get)
|
|
2444
|
+
|
|
2445
|
+
# Mock retry_request to call through, which will raise the 404 error
|
|
2446
|
+
async def mock_retry(func, *args, **kwargs):
|
|
2447
|
+
return await func(*args, **kwargs)
|
|
2448
|
+
|
|
2449
|
+
with patch(
|
|
2450
|
+
"proximl.utils.transfer.retry_request",
|
|
2451
|
+
side_effect=mock_retry,
|
|
2452
|
+
):
|
|
2453
|
+
with patch(
|
|
2454
|
+
"proximl.utils.transfer.ping_endpoint",
|
|
2455
|
+
new_callable=AsyncMock,
|
|
2456
|
+
):
|
|
2457
|
+
with raises(ClientResponseError):
|
|
2458
|
+
await specimen.download(
|
|
2459
|
+
"example.com", "token", tmpdir
|
|
2460
|
+
)
|
|
2461
|
+
|
|
2462
|
+
@mark.asyncio
|
|
2463
|
+
async def test_download_non_404_error_status(self):
|
|
2464
|
+
with tempfile.TemporaryDirectory() as tmpdir:
|
|
2465
|
+
with patch("aiohttp.ClientSession") as mock_session:
|
|
2466
|
+
mock_session_instance = AsyncMock()
|
|
2467
|
+
mock_session.return_value.__aenter__ = AsyncMock(
|
|
2468
|
+
return_value=mock_session_instance
|
|
2469
|
+
)
|
|
2470
|
+
mock_session.return_value.__aexit__ = AsyncMock(
|
|
2471
|
+
return_value=None
|
|
2472
|
+
)
|
|
2473
|
+
|
|
2474
|
+
def mock_get(*args, **kwargs):
|
|
2475
|
+
url = args[0] if args else kwargs.get("url", "")
|
|
2476
|
+
if "/info" in url or "/ping" in url:
|
|
2477
|
+
mock_resp = AsyncMock()
|
|
2478
|
+
mock_resp.status = 200
|
|
2479
|
+
mock_resp.json = AsyncMock(
|
|
2480
|
+
return_value={"archive": False}
|
|
2481
|
+
)
|
|
2482
|
+
mock_resp.text = AsyncMock(return_value="")
|
|
2483
|
+
mock_resp.__aenter__ = AsyncMock(
|
|
2484
|
+
return_value=mock_resp
|
|
2485
|
+
)
|
|
2486
|
+
mock_resp.__aexit__ = AsyncMock(return_value=None)
|
|
2487
|
+
return mock_resp
|
|
2488
|
+
else:
|
|
2489
|
+
# For /download endpoint, return awaitable that raises 500
|
|
2490
|
+
async def get_download_response():
|
|
2491
|
+
raise ClientResponseError(
|
|
2492
|
+
request_info=Mock(),
|
|
2493
|
+
history=(),
|
|
2494
|
+
status=500,
|
|
2495
|
+
message="Server error",
|
|
2496
|
+
)
|
|
2497
|
+
|
|
2498
|
+
return get_download_response()
|
|
2499
|
+
|
|
2500
|
+
mock_session_instance.get = Mock(side_effect=mock_get)
|
|
2501
|
+
|
|
2502
|
+
# Mock retry_request to call through, which will raise the 500 error
|
|
2503
|
+
async def mock_retry(func, *args, **kwargs):
|
|
2504
|
+
return await func(*args, **kwargs)
|
|
2505
|
+
|
|
2506
|
+
with patch(
|
|
2507
|
+
"proximl.utils.transfer.retry_request",
|
|
2508
|
+
side_effect=mock_retry,
|
|
2509
|
+
):
|
|
2510
|
+
with patch(
|
|
2511
|
+
"proximl.utils.transfer.ping_endpoint",
|
|
2512
|
+
new_callable=AsyncMock,
|
|
2513
|
+
):
|
|
2514
|
+
with raises(ClientResponseError):
|
|
2515
|
+
await specimen.download(
|
|
2516
|
+
"example.com", "token", tmpdir
|
|
2517
|
+
)
|
|
2518
|
+
|
|
2519
|
+
@mark.asyncio
|
|
2520
|
+
async def test_download_finalize_failure(self):
|
|
2521
|
+
with tempfile.TemporaryDirectory() as tmpdir:
|
|
2522
|
+
with patch("aiohttp.ClientSession") as mock_session:
|
|
2523
|
+
mock_session_instance = AsyncMock()
|
|
2524
|
+
mock_session.return_value.__aenter__ = AsyncMock(
|
|
2525
|
+
return_value=mock_session_instance
|
|
2526
|
+
)
|
|
2527
|
+
mock_session.return_value.__aexit__ = AsyncMock(
|
|
2528
|
+
return_value=None
|
|
2529
|
+
)
|
|
2530
|
+
|
|
2531
|
+
def mock_get(*args, **kwargs):
|
|
2532
|
+
url = args[0] if args else kwargs.get("url", "")
|
|
2533
|
+
if "/info" in url or "/ping" in url:
|
|
2534
|
+
mock_resp = AsyncMock()
|
|
2535
|
+
mock_resp.status = 200
|
|
2536
|
+
mock_resp.json = AsyncMock(
|
|
2537
|
+
return_value={"archive": False}
|
|
2538
|
+
)
|
|
2539
|
+
mock_resp.text = AsyncMock(return_value="")
|
|
2540
|
+
mock_resp.__aenter__ = AsyncMock(
|
|
2541
|
+
return_value=mock_resp
|
|
2542
|
+
)
|
|
2543
|
+
mock_resp.__aexit__ = AsyncMock(return_value=None)
|
|
2544
|
+
return mock_resp
|
|
2545
|
+
else:
|
|
2546
|
+
# For /download endpoint, return awaitable that resolves to response
|
|
2547
|
+
async def get_download_response():
|
|
2548
|
+
mock_resp = AsyncMock()
|
|
2549
|
+
mock_resp.status = 200
|
|
2550
|
+
mock_resp.headers = {
|
|
2551
|
+
"Content-Type": "application/x-tar"
|
|
2552
|
+
}
|
|
2553
|
+
|
|
2554
|
+
async def chunk_iter():
|
|
2555
|
+
yield b"tar data"
|
|
2556
|
+
yield b""
|
|
2557
|
+
|
|
2558
|
+
mock_resp.content.iter_chunked = (
|
|
2559
|
+
lambda size: chunk_iter()
|
|
2560
|
+
)
|
|
2561
|
+
mock_resp.close = Mock()
|
|
2562
|
+
return mock_resp
|
|
2563
|
+
|
|
2564
|
+
return get_download_response()
|
|
2565
|
+
|
|
2566
|
+
mock_session_instance.get = Mock(side_effect=mock_get)
|
|
2567
|
+
|
|
2568
|
+
# retry_request is called 3 times: _get_info, _download, _finalize
|
|
2569
|
+
async def mock_retry(func, *args, **kwargs):
|
|
2570
|
+
return await func(*args, **kwargs)
|
|
2571
|
+
|
|
2572
|
+
with patch(
|
|
2573
|
+
"proximl.utils.transfer.retry_request",
|
|
2574
|
+
side_effect=mock_retry,
|
|
2575
|
+
):
|
|
2576
|
+
with patch(
|
|
2577
|
+
"asyncio.create_subprocess_exec"
|
|
2578
|
+
) as mock_subprocess:
|
|
2579
|
+
mock_process = AsyncMock()
|
|
2580
|
+
mock_process.stdin = Mock()
|
|
2581
|
+
mock_process.stdin.write = Mock()
|
|
2582
|
+
mock_process.stdin.drain = AsyncMock()
|
|
2583
|
+
mock_process.stdin.close = Mock()
|
|
2584
|
+
mock_process.returncode = 0
|
|
2585
|
+
mock_process.wait = AsyncMock(return_value=0)
|
|
2586
|
+
mock_process.stderr.read = AsyncMock(return_value=b"")
|
|
2587
|
+
mock_subprocess.return_value = mock_process
|
|
2588
|
+
|
|
2589
|
+
mock_finalize_response = AsyncMock()
|
|
2590
|
+
mock_finalize_response.status = 500
|
|
2591
|
+
|
|
2592
|
+
mock_finalize_response.text = AsyncMock(
|
|
2593
|
+
return_value="Finalize error"
|
|
2594
|
+
)
|
|
2595
|
+
|
|
2596
|
+
mock_finalize_response.__aenter__ = AsyncMock(
|
|
2597
|
+
return_value=mock_finalize_response
|
|
2598
|
+
)
|
|
2599
|
+
|
|
2600
|
+
mock_finalize_response.__aexit__ = AsyncMock(
|
|
2601
|
+
return_value=None
|
|
2602
|
+
)
|
|
2603
|
+
|
|
2604
|
+
# session.post() should return something that is both awaitable and an async context manager
|
|
2605
|
+
class AwaitableContextManager:
|
|
2606
|
+
def __init__(self, return_value):
|
|
2607
|
+
self.return_value = return_value
|
|
2608
|
+
|
|
2609
|
+
def __await__(self):
|
|
2610
|
+
yield
|
|
2611
|
+
return self
|
|
2612
|
+
|
|
2613
|
+
async def __aenter__(self):
|
|
2614
|
+
return self.return_value
|
|
2615
|
+
|
|
2616
|
+
async def __aexit__(self, *args):
|
|
2617
|
+
return None
|
|
2618
|
+
|
|
2619
|
+
mock_post_context = AwaitableContextManager(
|
|
2620
|
+
mock_finalize_response
|
|
2621
|
+
)
|
|
2622
|
+
mock_session_instance.post = Mock(
|
|
2623
|
+
return_value=mock_post_context
|
|
2624
|
+
)
|
|
2625
|
+
|
|
2626
|
+
with patch(
|
|
2627
|
+
"proximl.utils.transfer.ping_endpoint",
|
|
2628
|
+
new_callable=AsyncMock,
|
|
2629
|
+
):
|
|
2630
|
+
with raises(
|
|
2631
|
+
ConnectionError, match="Finalize failed"
|
|
2632
|
+
):
|
|
2633
|
+
await specimen.download(
|
|
2634
|
+
"example.com", "token", tmpdir
|
|
2635
|
+
)
|
|
2636
|
+
|
|
2637
|
+
@mark.asyncio
|
|
2638
|
+
async def test_download_content_disposition_filename(self):
|
|
2639
|
+
with tempfile.TemporaryDirectory() as tmpdir:
|
|
2640
|
+
with patch("aiohttp.ClientSession") as mock_session:
|
|
2641
|
+
mock_session_instance = AsyncMock()
|
|
2642
|
+
mock_session.return_value.__aenter__ = AsyncMock(
|
|
2643
|
+
return_value=mock_session_instance
|
|
2644
|
+
)
|
|
2645
|
+
mock_session.return_value.__aexit__ = AsyncMock(
|
|
2646
|
+
return_value=None
|
|
2647
|
+
)
|
|
2648
|
+
|
|
2649
|
+
def mock_get(*args, **kwargs):
|
|
2650
|
+
url = args[0] if args else kwargs.get("url", "")
|
|
2651
|
+
if "/info" in url or "/ping" in url:
|
|
2652
|
+
mock_resp = AsyncMock()
|
|
2653
|
+
mock_resp.status = 200
|
|
2654
|
+
mock_resp.json = AsyncMock(
|
|
2655
|
+
return_value={"archive": True}
|
|
2656
|
+
)
|
|
2657
|
+
mock_resp.text = AsyncMock(return_value="")
|
|
2658
|
+
mock_resp.__aenter__ = AsyncMock(
|
|
2659
|
+
return_value=mock_resp
|
|
2660
|
+
)
|
|
2661
|
+
mock_resp.__aexit__ = AsyncMock(return_value=None)
|
|
2662
|
+
return mock_resp
|
|
2663
|
+
else:
|
|
2664
|
+
# For /download endpoint, return awaitable that resolves to response
|
|
2665
|
+
async def get_download_response():
|
|
2666
|
+
mock_resp = AsyncMock()
|
|
2667
|
+
mock_resp.status = 200
|
|
2668
|
+
mock_resp.headers = {
|
|
2669
|
+
"Content-Type": "application/zip",
|
|
2670
|
+
"Content-Disposition": 'attachment; filename="custom-name.zip"',
|
|
2671
|
+
}
|
|
2672
|
+
|
|
2673
|
+
async def chunk_iter():
|
|
2674
|
+
yield b"zip data"
|
|
2675
|
+
yield b""
|
|
2676
|
+
|
|
2677
|
+
mock_resp.content.iter_chunked = (
|
|
2678
|
+
lambda size: chunk_iter()
|
|
2679
|
+
)
|
|
2680
|
+
mock_resp.close = Mock()
|
|
2681
|
+
return mock_resp
|
|
2682
|
+
|
|
2683
|
+
return get_download_response()
|
|
2684
|
+
|
|
2685
|
+
mock_session_instance.get = Mock(side_effect=mock_get)
|
|
2686
|
+
|
|
2687
|
+
# Mock retry_request to actually call the function passed to it
|
|
2688
|
+
async def mock_retry(func, *args, **kwargs):
|
|
2689
|
+
return await func(*args, **kwargs)
|
|
2690
|
+
|
|
2691
|
+
with patch(
|
|
2692
|
+
"proximl.utils.transfer.retry_request",
|
|
2693
|
+
side_effect=mock_retry,
|
|
2694
|
+
):
|
|
2695
|
+
mock_file_context = AsyncMock()
|
|
2696
|
+
mock_file_context.__aenter__ = AsyncMock(
|
|
2697
|
+
return_value=AsyncMock()
|
|
2698
|
+
)
|
|
2699
|
+
mock_file_context.__aexit__ = AsyncMock(return_value=None)
|
|
2700
|
+
mock_file_context.__aenter__.return_value.write = (
|
|
2701
|
+
AsyncMock()
|
|
2702
|
+
)
|
|
2703
|
+
with patch(
|
|
2704
|
+
"aiofiles.open", return_value=mock_file_context
|
|
2705
|
+
) as mock_file:
|
|
2706
|
+
mock_finalize_response = AsyncMock()
|
|
2707
|
+
mock_finalize_response.status = 200
|
|
2708
|
+
mock_finalize_response.json = AsyncMock(
|
|
2709
|
+
return_value={"status": "ok"}
|
|
2710
|
+
)
|
|
2711
|
+
|
|
2712
|
+
mock_finalize_response.__aenter__ = AsyncMock(
|
|
2713
|
+
return_value=mock_finalize_response
|
|
2714
|
+
)
|
|
2715
|
+
|
|
2716
|
+
mock_finalize_response.__aexit__ = AsyncMock(
|
|
2717
|
+
return_value=None
|
|
2718
|
+
)
|
|
2719
|
+
|
|
2720
|
+
# session.post() should return something that is both awaitable and an async context manager
|
|
2721
|
+
class AwaitableContextManager:
|
|
2722
|
+
def __init__(self, return_value):
|
|
2723
|
+
self.return_value = return_value
|
|
2724
|
+
|
|
2725
|
+
def __await__(self):
|
|
2726
|
+
yield
|
|
2727
|
+
return self
|
|
2728
|
+
|
|
2729
|
+
async def __aenter__(self):
|
|
2730
|
+
return self.return_value
|
|
2731
|
+
|
|
2732
|
+
async def __aexit__(self, *args):
|
|
2733
|
+
return None
|
|
2734
|
+
|
|
2735
|
+
mock_post_context = AwaitableContextManager(
|
|
2736
|
+
mock_finalize_response
|
|
2737
|
+
)
|
|
2738
|
+
mock_session_instance.post = Mock(
|
|
2739
|
+
return_value=mock_post_context
|
|
2740
|
+
)
|
|
2741
|
+
|
|
2742
|
+
with patch(
|
|
2743
|
+
"proximl.utils.transfer.ping_endpoint",
|
|
2744
|
+
new_callable=AsyncMock,
|
|
2745
|
+
):
|
|
2746
|
+
await specimen.download(
|
|
2747
|
+
"example.com", "token", tmpdir
|
|
2748
|
+
)
|
|
2749
|
+
|
|
2750
|
+
@mark.asyncio
|
|
2751
|
+
async def test_download_content_disposition_filename_no_quotes(self):
|
|
2752
|
+
with tempfile.TemporaryDirectory() as tmpdir:
|
|
2753
|
+
with patch("aiohttp.ClientSession") as mock_session:
|
|
2754
|
+
mock_session_instance = AsyncMock()
|
|
2755
|
+
mock_session.return_value.__aenter__ = AsyncMock(
|
|
2756
|
+
return_value=mock_session_instance
|
|
2757
|
+
)
|
|
2758
|
+
mock_session.return_value.__aexit__ = AsyncMock(
|
|
2759
|
+
return_value=None
|
|
2760
|
+
)
|
|
2761
|
+
|
|
2762
|
+
def mock_get(*args, **kwargs):
|
|
2763
|
+
url = args[0] if args else kwargs.get("url", "")
|
|
2764
|
+
# Handle /ping endpoint for ping_endpoint
|
|
2765
|
+
if "/ping" in url:
|
|
2766
|
+
mock_resp = AsyncMock()
|
|
2767
|
+
mock_resp.status = 200
|
|
2768
|
+
mock_resp.__aenter__ = AsyncMock(
|
|
2769
|
+
return_value=mock_resp
|
|
2770
|
+
)
|
|
2771
|
+
mock_resp.__aexit__ = AsyncMock(return_value=None)
|
|
2772
|
+
return mock_resp
|
|
2773
|
+
# Handle /info endpoint
|
|
2774
|
+
elif "/info" in url:
|
|
2775
|
+
mock_resp = AsyncMock()
|
|
2776
|
+
mock_resp.status = 200
|
|
2777
|
+
mock_resp.json = AsyncMock(
|
|
2778
|
+
return_value={"archive": True}
|
|
2779
|
+
)
|
|
2780
|
+
mock_resp.text = AsyncMock(return_value="")
|
|
2781
|
+
mock_resp.__aenter__ = AsyncMock(
|
|
2782
|
+
return_value=mock_resp
|
|
2783
|
+
)
|
|
2784
|
+
mock_resp.__aexit__ = AsyncMock(return_value=None)
|
|
2785
|
+
return mock_resp
|
|
2786
|
+
else:
|
|
2787
|
+
# For /download endpoint, return awaitable that resolves to response
|
|
2788
|
+
async def get_download_response():
|
|
2789
|
+
mock_resp = AsyncMock()
|
|
2790
|
+
mock_resp.status = 200
|
|
2791
|
+
mock_resp.headers = {
|
|
2792
|
+
"Content-Type": "application/zip",
|
|
2793
|
+
"Content-Disposition": "attachment; filename=custom-name.zip",
|
|
2794
|
+
}
|
|
2795
|
+
|
|
2796
|
+
async def chunk_iter():
|
|
2797
|
+
yield b"zip data"
|
|
2798
|
+
yield b""
|
|
2799
|
+
|
|
2800
|
+
mock_resp.content.iter_chunked = (
|
|
2801
|
+
lambda size: chunk_iter()
|
|
2802
|
+
)
|
|
2803
|
+
mock_resp.close = Mock()
|
|
2804
|
+
return mock_resp
|
|
2805
|
+
|
|
2806
|
+
return get_download_response()
|
|
2807
|
+
|
|
2808
|
+
mock_session_instance.get = Mock(side_effect=mock_get)
|
|
2809
|
+
|
|
2810
|
+
# Mock retry_request to actually call the function passed to it
|
|
2811
|
+
async def mock_retry(func, *args, **kwargs):
|
|
2812
|
+
return await func(*args, **kwargs)
|
|
2813
|
+
|
|
2814
|
+
with patch(
|
|
2815
|
+
"proximl.utils.transfer.retry_request",
|
|
2816
|
+
side_effect=mock_retry,
|
|
2817
|
+
):
|
|
2818
|
+
with patch(
|
|
2819
|
+
"proximl.utils.transfer.ping_endpoint",
|
|
2820
|
+
new_callable=AsyncMock,
|
|
2821
|
+
):
|
|
2822
|
+
mock_file_context = AsyncMock()
|
|
2823
|
+
mock_file_context.__aenter__ = AsyncMock(
|
|
2824
|
+
return_value=AsyncMock()
|
|
2825
|
+
)
|
|
2826
|
+
mock_file_context.__aexit__ = AsyncMock(
|
|
2827
|
+
return_value=None
|
|
2828
|
+
)
|
|
2829
|
+
mock_file_context.__aenter__.return_value.write = (
|
|
2830
|
+
AsyncMock()
|
|
2831
|
+
)
|
|
2832
|
+
with patch(
|
|
2833
|
+
"aiofiles.open", return_value=mock_file_context
|
|
2834
|
+
) as mock_file:
|
|
2835
|
+
mock_finalize_response = AsyncMock()
|
|
2836
|
+
mock_finalize_response.status = 200
|
|
2837
|
+
mock_finalize_response.json = AsyncMock(
|
|
2838
|
+
return_value={"status": "ok"}
|
|
2839
|
+
)
|
|
2840
|
+
|
|
2841
|
+
mock_finalize_response.__aenter__ = AsyncMock(
|
|
2842
|
+
return_value=mock_finalize_response
|
|
2843
|
+
)
|
|
2844
|
+
|
|
2845
|
+
mock_finalize_response.__aexit__ = AsyncMock(
|
|
2846
|
+
return_value=None
|
|
2847
|
+
)
|
|
2848
|
+
|
|
2849
|
+
# session.post() should return something that is both awaitable and an async context manager
|
|
2850
|
+
class AwaitableContextManager:
|
|
2851
|
+
def __init__(self, return_value):
|
|
2852
|
+
self.return_value = return_value
|
|
2853
|
+
|
|
2854
|
+
def __await__(self):
|
|
2855
|
+
yield
|
|
2856
|
+
return self
|
|
2857
|
+
|
|
2858
|
+
async def __aenter__(self):
|
|
2859
|
+
return self.return_value
|
|
2860
|
+
|
|
2861
|
+
async def __aexit__(self, *args):
|
|
2862
|
+
return None
|
|
2863
|
+
|
|
2864
|
+
mock_post_context = AwaitableContextManager(
|
|
2865
|
+
mock_finalize_response
|
|
2866
|
+
)
|
|
2867
|
+
mock_session_instance.post = Mock(
|
|
2868
|
+
return_value=mock_post_context
|
|
2869
|
+
)
|
|
2870
|
+
|
|
2871
|
+
with patch(
|
|
2872
|
+
"proximl.utils.transfer.ping_endpoint",
|
|
2873
|
+
new_callable=AsyncMock,
|
|
2874
|
+
):
|
|
2875
|
+
await specimen.download(
|
|
2876
|
+
"example.com", "token", tmpdir
|
|
2877
|
+
)
|
|
2878
|
+
|
|
2879
|
+
@mark.asyncio
|
|
2880
|
+
async def test_download_no_content_disposition_fallback(self):
|
|
2881
|
+
with tempfile.TemporaryDirectory() as tmpdir:
|
|
2882
|
+
with patch("aiohttp.ClientSession") as mock_session:
|
|
2883
|
+
mock_session_instance = AsyncMock()
|
|
2884
|
+
mock_session.return_value.__aenter__ = AsyncMock(
|
|
2885
|
+
return_value=mock_session_instance
|
|
2886
|
+
)
|
|
2887
|
+
mock_session.return_value.__aexit__ = AsyncMock(
|
|
2888
|
+
return_value=None
|
|
2889
|
+
)
|
|
2890
|
+
|
|
2891
|
+
def mock_get(*args, **kwargs):
|
|
2892
|
+
url = args[0] if args else kwargs.get("url", "")
|
|
2893
|
+
if "/info" in url or "/ping" in url:
|
|
2894
|
+
mock_resp = AsyncMock()
|
|
2895
|
+
mock_resp.status = 200
|
|
2896
|
+
mock_resp.json = AsyncMock(
|
|
2897
|
+
return_value={"archive": True}
|
|
2898
|
+
)
|
|
2899
|
+
mock_resp.text = AsyncMock(return_value="")
|
|
2900
|
+
mock_resp.__aenter__ = AsyncMock(
|
|
2901
|
+
return_value=mock_resp
|
|
2902
|
+
)
|
|
2903
|
+
mock_resp.__aexit__ = AsyncMock(return_value=None)
|
|
2904
|
+
return mock_resp
|
|
2905
|
+
else:
|
|
2906
|
+
# For /download endpoint, return awaitable that resolves to response
|
|
2907
|
+
async def get_download_response():
|
|
2908
|
+
mock_resp = AsyncMock()
|
|
2909
|
+
mock_resp.status = 200
|
|
2910
|
+
mock_resp.headers = {
|
|
2911
|
+
"Content-Type": "application/zip",
|
|
2912
|
+
}
|
|
2913
|
+
|
|
2914
|
+
async def chunk_iter():
|
|
2915
|
+
yield b"zip data"
|
|
2916
|
+
yield b""
|
|
2917
|
+
|
|
2918
|
+
mock_resp.content.iter_chunked = (
|
|
2919
|
+
lambda size: chunk_iter()
|
|
2920
|
+
)
|
|
2921
|
+
mock_resp.close = Mock()
|
|
2922
|
+
return mock_resp
|
|
2923
|
+
|
|
2924
|
+
return get_download_response()
|
|
2925
|
+
|
|
2926
|
+
mock_session_instance.get = Mock(side_effect=mock_get)
|
|
2927
|
+
|
|
2928
|
+
# Mock retry_request to actually call the function passed to it
|
|
2929
|
+
async def mock_retry(func, *args, **kwargs):
|
|
2930
|
+
return await func(*args, **kwargs)
|
|
2931
|
+
|
|
2932
|
+
with patch(
|
|
2933
|
+
"proximl.utils.transfer.retry_request",
|
|
2934
|
+
side_effect=mock_retry,
|
|
2935
|
+
):
|
|
2936
|
+
mock_file_context = AsyncMock()
|
|
2937
|
+
mock_file_context.__aenter__ = AsyncMock(
|
|
2938
|
+
return_value=AsyncMock()
|
|
2939
|
+
)
|
|
2940
|
+
mock_file_context.__aexit__ = AsyncMock(return_value=None)
|
|
2941
|
+
mock_file_context.__aenter__.return_value.write = (
|
|
2942
|
+
AsyncMock()
|
|
2943
|
+
)
|
|
2944
|
+
with patch(
|
|
2945
|
+
"aiofiles.open", return_value=mock_file_context
|
|
2946
|
+
) as mock_file:
|
|
2947
|
+
mock_finalize_response = AsyncMock()
|
|
2948
|
+
mock_finalize_response.status = 200
|
|
2949
|
+
mock_finalize_response.json = AsyncMock(
|
|
2950
|
+
return_value={"status": "ok"}
|
|
2951
|
+
)
|
|
2952
|
+
|
|
2953
|
+
mock_finalize_response.__aenter__ = AsyncMock(
|
|
2954
|
+
return_value=mock_finalize_response
|
|
2955
|
+
)
|
|
2956
|
+
|
|
2957
|
+
mock_finalize_response.__aexit__ = AsyncMock(
|
|
2958
|
+
return_value=None
|
|
2959
|
+
)
|
|
2960
|
+
|
|
2961
|
+
# session.post() should return something that is both awaitable and an async context manager
|
|
2962
|
+
class AwaitableContextManager:
|
|
2963
|
+
def __init__(self, return_value):
|
|
2964
|
+
self.return_value = return_value
|
|
2965
|
+
|
|
2966
|
+
def __await__(self):
|
|
2967
|
+
yield
|
|
2968
|
+
return self
|
|
2969
|
+
|
|
2970
|
+
async def __aenter__(self):
|
|
2971
|
+
return self.return_value
|
|
2972
|
+
|
|
2973
|
+
async def __aexit__(self, *args):
|
|
2974
|
+
return None
|
|
2975
|
+
|
|
2976
|
+
mock_post_context = AwaitableContextManager(
|
|
2977
|
+
mock_finalize_response
|
|
2978
|
+
)
|
|
2979
|
+
mock_session_instance.post = Mock(
|
|
2980
|
+
return_value=mock_post_context
|
|
2981
|
+
)
|
|
2982
|
+
|
|
2983
|
+
with patch(
|
|
2984
|
+
"proximl.utils.transfer.ping_endpoint",
|
|
2985
|
+
new_callable=AsyncMock,
|
|
2986
|
+
):
|
|
2987
|
+
await specimen.download(
|
|
2988
|
+
"example.com", "token", tmpdir
|
|
2989
|
+
)
|
|
2990
|
+
|
|
2991
|
+
@mark.asyncio
|
|
2992
|
+
async def test_download_filename_no_zip_extension(self):
|
|
2993
|
+
with tempfile.TemporaryDirectory() as tmpdir:
|
|
2994
|
+
with patch("aiohttp.ClientSession") as mock_session:
|
|
2995
|
+
mock_session_instance = AsyncMock()
|
|
2996
|
+
mock_session.return_value.__aenter__ = AsyncMock(
|
|
2997
|
+
return_value=mock_session_instance
|
|
2998
|
+
)
|
|
2999
|
+
mock_session.return_value.__aexit__ = AsyncMock(
|
|
3000
|
+
return_value=None
|
|
3001
|
+
)
|
|
3002
|
+
|
|
3003
|
+
def mock_get(*args, **kwargs):
|
|
3004
|
+
url = args[0] if args else kwargs.get("url", "")
|
|
3005
|
+
# Handle /ping endpoint for ping_endpoint
|
|
3006
|
+
if "/ping" in url:
|
|
3007
|
+
mock_resp = AsyncMock()
|
|
3008
|
+
mock_resp.status = 200
|
|
3009
|
+
mock_resp.__aenter__ = AsyncMock(
|
|
3010
|
+
return_value=mock_resp
|
|
3011
|
+
)
|
|
3012
|
+
mock_resp.__aexit__ = AsyncMock(return_value=None)
|
|
3013
|
+
return mock_resp
|
|
3014
|
+
# Handle /info endpoint
|
|
3015
|
+
elif "/info" in url:
|
|
3016
|
+
mock_resp = AsyncMock()
|
|
3017
|
+
mock_resp.status = 200
|
|
3018
|
+
mock_resp.json = AsyncMock(
|
|
3019
|
+
return_value={"archive": True}
|
|
3020
|
+
)
|
|
3021
|
+
mock_resp.text = AsyncMock(return_value="")
|
|
3022
|
+
mock_resp.__aenter__ = AsyncMock(
|
|
3023
|
+
return_value=mock_resp
|
|
3024
|
+
)
|
|
3025
|
+
mock_resp.__aexit__ = AsyncMock(return_value=None)
|
|
3026
|
+
return mock_resp
|
|
3027
|
+
else:
|
|
3028
|
+
# For /download endpoint, return awaitable that resolves to response
|
|
3029
|
+
async def get_download_response():
|
|
3030
|
+
mock_resp = AsyncMock()
|
|
3031
|
+
mock_resp.status = 200
|
|
3032
|
+
mock_resp.headers = {
|
|
3033
|
+
"Content-Type": "application/zip",
|
|
3034
|
+
"Content-Disposition": 'attachment; filename="custom-name"',
|
|
3035
|
+
}
|
|
3036
|
+
|
|
3037
|
+
async def chunk_iter():
|
|
3038
|
+
yield b"zip data"
|
|
3039
|
+
yield b""
|
|
3040
|
+
|
|
3041
|
+
mock_resp.content.iter_chunked = (
|
|
3042
|
+
lambda size: chunk_iter()
|
|
3043
|
+
)
|
|
3044
|
+
mock_resp.close = Mock()
|
|
3045
|
+
return mock_resp
|
|
3046
|
+
|
|
3047
|
+
return get_download_response()
|
|
3048
|
+
|
|
3049
|
+
mock_session_instance.get = Mock(side_effect=mock_get)
|
|
3050
|
+
|
|
3051
|
+
# Mock retry_request to actually call the function passed to it
|
|
3052
|
+
async def mock_retry(func, *args, **kwargs):
|
|
3053
|
+
return await func(*args, **kwargs)
|
|
3054
|
+
|
|
3055
|
+
with patch(
|
|
3056
|
+
"proximl.utils.transfer.retry_request",
|
|
3057
|
+
side_effect=mock_retry,
|
|
3058
|
+
):
|
|
3059
|
+
with patch(
|
|
3060
|
+
"proximl.utils.transfer.ping_endpoint",
|
|
3061
|
+
new_callable=AsyncMock,
|
|
3062
|
+
):
|
|
3063
|
+
mock_file_context = AsyncMock()
|
|
3064
|
+
mock_file_context.__aenter__ = AsyncMock(
|
|
3065
|
+
return_value=AsyncMock()
|
|
3066
|
+
)
|
|
3067
|
+
mock_file_context.__aexit__ = AsyncMock(
|
|
3068
|
+
return_value=None
|
|
3069
|
+
)
|
|
3070
|
+
mock_file_context.__aenter__.return_value.write = (
|
|
3071
|
+
AsyncMock()
|
|
3072
|
+
)
|
|
3073
|
+
with patch(
|
|
3074
|
+
"aiofiles.open", return_value=mock_file_context
|
|
3075
|
+
) as mock_file:
|
|
3076
|
+
mock_finalize_response = AsyncMock()
|
|
3077
|
+
mock_finalize_response.status = 200
|
|
3078
|
+
mock_finalize_response.json = AsyncMock(
|
|
3079
|
+
return_value={"status": "ok"}
|
|
3080
|
+
)
|
|
3081
|
+
|
|
3082
|
+
mock_finalize_response.__aenter__ = AsyncMock(
|
|
3083
|
+
return_value=mock_finalize_response
|
|
3084
|
+
)
|
|
3085
|
+
|
|
3086
|
+
mock_finalize_response.__aexit__ = AsyncMock(
|
|
3087
|
+
return_value=None
|
|
3088
|
+
)
|
|
3089
|
+
|
|
3090
|
+
# session.post() should return something that is both awaitable and an async context manager
|
|
3091
|
+
class AwaitableContextManager:
|
|
3092
|
+
def __init__(self, return_value):
|
|
3093
|
+
self.return_value = return_value
|
|
3094
|
+
|
|
3095
|
+
def __await__(self):
|
|
3096
|
+
yield
|
|
3097
|
+
return self
|
|
3098
|
+
|
|
3099
|
+
async def __aenter__(self):
|
|
3100
|
+
return self.return_value
|
|
3101
|
+
|
|
3102
|
+
async def __aexit__(self, *args):
|
|
3103
|
+
return None
|
|
3104
|
+
|
|
3105
|
+
mock_post_context = AwaitableContextManager(
|
|
3106
|
+
mock_finalize_response
|
|
3107
|
+
)
|
|
3108
|
+
mock_session_instance.post = Mock(
|
|
3109
|
+
return_value=mock_post_context
|
|
3110
|
+
)
|
|
3111
|
+
|
|
3112
|
+
with patch(
|
|
3113
|
+
"proximl.utils.transfer.ping_endpoint",
|
|
3114
|
+
new_callable=AsyncMock,
|
|
3115
|
+
):
|
|
3116
|
+
await specimen.download(
|
|
3117
|
+
"example.com", "token", tmpdir
|
|
3118
|
+
)
|
|
3119
|
+
|
|
3120
|
+
@mark.asyncio
|
|
3121
|
+
async def test_download_multiple_chunks(self):
|
|
3122
|
+
with tempfile.TemporaryDirectory() as tmpdir:
|
|
3123
|
+
with patch("aiohttp.ClientSession") as mock_session:
|
|
3124
|
+
mock_session_instance = AsyncMock()
|
|
3125
|
+
mock_session.return_value.__aenter__ = AsyncMock(
|
|
3126
|
+
return_value=mock_session_instance
|
|
3127
|
+
)
|
|
3128
|
+
mock_session.return_value.__aexit__ = AsyncMock(
|
|
3129
|
+
return_value=None
|
|
3130
|
+
)
|
|
3131
|
+
|
|
3132
|
+
def mock_get(*args, **kwargs):
|
|
3133
|
+
url = args[0] if args else kwargs.get("url", "")
|
|
3134
|
+
# Handle /ping endpoint for ping_endpoint
|
|
3135
|
+
if "/ping" in url:
|
|
3136
|
+
mock_resp = AsyncMock()
|
|
3137
|
+
mock_resp.status = 200
|
|
3138
|
+
mock_resp.__aenter__ = AsyncMock(
|
|
3139
|
+
return_value=mock_resp
|
|
3140
|
+
)
|
|
3141
|
+
mock_resp.__aexit__ = AsyncMock(return_value=None)
|
|
3142
|
+
return mock_resp
|
|
3143
|
+
# Handle /info endpoint
|
|
3144
|
+
elif "/info" in url:
|
|
3145
|
+
mock_resp = AsyncMock()
|
|
3146
|
+
mock_resp.status = 200
|
|
3147
|
+
mock_resp.json = AsyncMock(
|
|
3148
|
+
return_value={"archive": False}
|
|
3149
|
+
)
|
|
3150
|
+
mock_resp.text = AsyncMock(return_value="")
|
|
3151
|
+
mock_resp.__aenter__ = AsyncMock(
|
|
3152
|
+
return_value=mock_resp
|
|
3153
|
+
)
|
|
3154
|
+
mock_resp.__aexit__ = AsyncMock(return_value=None)
|
|
3155
|
+
return mock_resp
|
|
3156
|
+
else:
|
|
3157
|
+
# For /download endpoint, return awaitable that resolves to response
|
|
3158
|
+
async def get_download_response():
|
|
3159
|
+
mock_resp = AsyncMock()
|
|
3160
|
+
mock_resp.status = 200
|
|
3161
|
+
mock_resp.headers = {
|
|
3162
|
+
"Content-Type": "application/x-tar"
|
|
3163
|
+
}
|
|
3164
|
+
|
|
3165
|
+
# Simulate multiple chunks - iter_chunked should return an async iterator
|
|
3166
|
+
async def chunk_iter():
|
|
3167
|
+
yield b"chunk1"
|
|
3168
|
+
yield b"chunk2"
|
|
3169
|
+
yield b""
|
|
3170
|
+
|
|
3171
|
+
mock_resp.content.iter_chunked = (
|
|
3172
|
+
lambda size: chunk_iter()
|
|
3173
|
+
)
|
|
3174
|
+
mock_resp.close = Mock()
|
|
3175
|
+
return mock_resp
|
|
3176
|
+
|
|
3177
|
+
return get_download_response()
|
|
3178
|
+
|
|
3179
|
+
mock_session_instance.get = Mock(side_effect=mock_get)
|
|
3180
|
+
|
|
3181
|
+
# retry_request is called 3 times: _get_info, _download, _finalize
|
|
3182
|
+
async def mock_retry(func, *args, **kwargs):
|
|
3183
|
+
return await func(*args, **kwargs)
|
|
3184
|
+
|
|
3185
|
+
with patch(
|
|
3186
|
+
"proximl.utils.transfer.retry_request",
|
|
3187
|
+
side_effect=mock_retry,
|
|
3188
|
+
):
|
|
3189
|
+
with patch(
|
|
3190
|
+
"proximl.utils.transfer.ping_endpoint",
|
|
3191
|
+
new_callable=AsyncMock,
|
|
3192
|
+
):
|
|
3193
|
+
with patch(
|
|
3194
|
+
"asyncio.create_subprocess_exec"
|
|
3195
|
+
) as mock_subprocess:
|
|
3196
|
+
mock_process = AsyncMock()
|
|
3197
|
+
mock_process.stdin = Mock()
|
|
3198
|
+
mock_process.stdin.write = Mock()
|
|
3199
|
+
mock_process.stdin.drain = AsyncMock()
|
|
3200
|
+
mock_process.stdin.close = Mock()
|
|
3201
|
+
mock_process.returncode = 0
|
|
3202
|
+
mock_process.wait = AsyncMock(return_value=0)
|
|
3203
|
+
mock_process.stderr.read = AsyncMock(
|
|
3204
|
+
return_value=b""
|
|
3205
|
+
)
|
|
3206
|
+
|
|
3207
|
+
# create_subprocess_exec is async, so return a coroutine
|
|
3208
|
+
async def mock_create_subprocess_exec(
|
|
3209
|
+
*args, **kwargs
|
|
3210
|
+
):
|
|
3211
|
+
return mock_process
|
|
3212
|
+
|
|
3213
|
+
mock_subprocess.side_effect = (
|
|
3214
|
+
mock_create_subprocess_exec
|
|
3215
|
+
)
|
|
3216
|
+
|
|
3217
|
+
mock_finalize_response = AsyncMock()
|
|
3218
|
+
mock_finalize_response.status = 200
|
|
3219
|
+
mock_finalize_response.json = AsyncMock(
|
|
3220
|
+
return_value={"status": "ok"}
|
|
3221
|
+
)
|
|
3222
|
+
mock_finalize_response.__aenter__ = AsyncMock(
|
|
3223
|
+
return_value=mock_finalize_response
|
|
3224
|
+
)
|
|
3225
|
+
mock_finalize_response.__aexit__ = AsyncMock(
|
|
3226
|
+
return_value=None
|
|
3227
|
+
)
|
|
3228
|
+
|
|
3229
|
+
# session.post() should return something that is both awaitable and an async context manager
|
|
3230
|
+
class AwaitableContextManager:
|
|
3231
|
+
def __init__(self, return_value):
|
|
3232
|
+
self.return_value = return_value
|
|
3233
|
+
|
|
3234
|
+
def __await__(self):
|
|
3235
|
+
yield
|
|
3236
|
+
return self
|
|
3237
|
+
|
|
3238
|
+
async def __aenter__(self):
|
|
3239
|
+
return self.return_value
|
|
3240
|
+
|
|
3241
|
+
async def __aexit__(self, *args):
|
|
3242
|
+
return None
|
|
3243
|
+
|
|
3244
|
+
mock_post_context = AwaitableContextManager(
|
|
3245
|
+
mock_finalize_response
|
|
3246
|
+
)
|
|
3247
|
+
mock_session_instance.post = Mock(
|
|
3248
|
+
return_value=mock_post_context
|
|
3249
|
+
)
|
|
3250
|
+
|
|
3251
|
+
await specimen.download(
|
|
3252
|
+
"example.com", "token", tmpdir
|
|
3253
|
+
)
|
|
3254
|
+
# Verify stdin.write was called for each chunk (chunk1, chunk2, and empty)
|
|
3255
|
+
# The empty chunk at the end also triggers a write
|
|
3256
|
+
assert mock_process.stdin.write.call_count >= 2
|
|
3257
|
+
|
|
3258
|
+
@mark.asyncio
|
|
3259
|
+
async def test_upload_chunk_retry_status_504(self):
|
|
3260
|
+
"""Test upload_chunk retry on 504 status - this tests lines 106-113"""
|
|
3261
|
+
session = AsyncMock()
|
|
3262
|
+
response = AsyncMock()
|
|
3263
|
+
response.status = 504
|
|
3264
|
+
response.text = AsyncMock(return_value="Gateway Timeout")
|
|
3265
|
+
response.request_info = Mock()
|
|
3266
|
+
response.history = ()
|
|
3267
|
+
session.put = AsyncMock(return_value=response.__aenter__())
|
|
3268
|
+
response.__aenter__ = AsyncMock(return_value=response)
|
|
3269
|
+
response.__aexit__ = AsyncMock(return_value=None)
|
|
3270
|
+
|
|
3271
|
+
async def _upload(*args, **kwargs):
|
|
3272
|
+
async with session.put(
|
|
3273
|
+
f"https://example.com/upload",
|
|
3274
|
+
headers={},
|
|
3275
|
+
data=b"data",
|
|
3276
|
+
timeout=30,
|
|
3277
|
+
) as resp:
|
|
3278
|
+
if resp.status == 504:
|
|
3279
|
+
text = await resp.text()
|
|
3280
|
+
raise ClientResponseError(
|
|
3281
|
+
request_info=resp.request_info,
|
|
3282
|
+
history=resp.history,
|
|
3283
|
+
status=resp.status,
|
|
3284
|
+
message=text,
|
|
3285
|
+
)
|
|
3286
|
+
|
|
3287
|
+
# Mock the actual upload_chunk behavior with retry
|
|
3288
|
+
with patch("proximl.utils.transfer.retry_request") as mock_retry:
|
|
3289
|
+
# Simulate retry_request calling _upload which raises 504
|
|
3290
|
+
# The retry_request will retry, but we'll make it fail after max retries
|
|
3291
|
+
mock_retry.side_effect = ClientResponseError(
|
|
3292
|
+
request_info=Mock(),
|
|
3293
|
+
history=(),
|
|
3294
|
+
status=504,
|
|
3295
|
+
message="Gateway Timeout",
|
|
3296
|
+
)
|
|
3297
|
+
with patch("asyncio.sleep", new_callable=AsyncMock):
|
|
3298
|
+
# This should trigger the retry logic, but we'll let it fail after max retries
|
|
3299
|
+
with raises(ClientResponseError):
|
|
3300
|
+
await specimen.upload_chunk(
|
|
3301
|
+
session,
|
|
3302
|
+
"https://example.com",
|
|
3303
|
+
"token",
|
|
3304
|
+
100,
|
|
3305
|
+
b"data",
|
|
3306
|
+
0,
|
|
3307
|
+
)
|
|
3308
|
+
|
|
3309
|
+
@mark.asyncio
|
|
3310
|
+
async def test_upload_finalize_success_logging(self):
|
|
3311
|
+
"""Test upload finalize success logging"""
|
|
3312
|
+
with tempfile.NamedTemporaryFile() as tmp:
|
|
3313
|
+
tmp.write(b"test content")
|
|
3314
|
+
tmp.flush()
|
|
3315
|
+
|
|
3316
|
+
with patch(
|
|
3317
|
+
"proximl.utils.transfer.ping_endpoint", new_callable=AsyncMock
|
|
3318
|
+
):
|
|
3319
|
+
with patch(
|
|
3320
|
+
"asyncio.create_subprocess_exec"
|
|
3321
|
+
) as mock_subprocess:
|
|
3322
|
+
mock_process = AsyncMock()
|
|
3323
|
+
mock_process.stdout.read = AsyncMock(
|
|
3324
|
+
side_effect=[b"data", b""]
|
|
3325
|
+
)
|
|
3326
|
+
mock_process.returncode = 0
|
|
3327
|
+
mock_process.wait = AsyncMock(return_value=0)
|
|
3328
|
+
mock_process.stderr.read = AsyncMock(return_value=b"")
|
|
3329
|
+
mock_subprocess.return_value = mock_process
|
|
3330
|
+
|
|
3331
|
+
with patch("aiohttp.ClientSession") as mock_session:
|
|
3332
|
+
mock_session_instance = AsyncMock()
|
|
3333
|
+
mock_session.return_value.__aenter__ = AsyncMock(
|
|
3334
|
+
return_value=mock_session_instance
|
|
3335
|
+
)
|
|
3336
|
+
mock_session.return_value.__aexit__ = AsyncMock(
|
|
3337
|
+
return_value=None
|
|
3338
|
+
)
|
|
3339
|
+
|
|
3340
|
+
with patch(
|
|
3341
|
+
"proximl.utils.transfer.upload_chunk",
|
|
3342
|
+
new_callable=AsyncMock,
|
|
3343
|
+
):
|
|
3344
|
+
mock_finalize_response = AsyncMock()
|
|
3345
|
+
mock_finalize_response.status = 200
|
|
3346
|
+
mock_finalize_response.json = AsyncMock(
|
|
3347
|
+
return_value={"status": "ok", "hash": "abc123"}
|
|
3348
|
+
)
|
|
3349
|
+
mock_finalize_response.__aenter__ = AsyncMock(
|
|
3350
|
+
return_value=mock_finalize_response
|
|
3351
|
+
)
|
|
3352
|
+
mock_finalize_response.__aexit__ = AsyncMock(
|
|
3353
|
+
return_value=None
|
|
3354
|
+
)
|
|
3355
|
+
|
|
3356
|
+
# session.post() should return something that is both awaitable and an async context manager
|
|
3357
|
+
class AwaitableContextManager:
|
|
3358
|
+
def __init__(self, return_value):
|
|
3359
|
+
self.return_value = return_value
|
|
3360
|
+
|
|
3361
|
+
def __await__(self):
|
|
3362
|
+
yield
|
|
3363
|
+
return self
|
|
3364
|
+
|
|
3365
|
+
async def __aenter__(self):
|
|
3366
|
+
return self.return_value
|
|
3367
|
+
|
|
3368
|
+
async def __aexit__(self, *args):
|
|
3369
|
+
return None
|
|
3370
|
+
|
|
3371
|
+
mock_post_context = AwaitableContextManager(
|
|
3372
|
+
mock_finalize_response
|
|
3373
|
+
)
|
|
3374
|
+
mock_session_instance.post = Mock(
|
|
3375
|
+
return_value=mock_post_context
|
|
3376
|
+
)
|
|
3377
|
+
|
|
3378
|
+
with patch("logging.debug") as mock_log:
|
|
3379
|
+
await specimen.upload(
|
|
3380
|
+
"example.com", "token", tmp.name
|
|
3381
|
+
)
|
|
3382
|
+
# Verify logging.debug was called for finalize
|
|
3383
|
+
mock_log.assert_called()
|
|
3384
|
+
|
|
3385
|
+
@mark.asyncio
|
|
3386
|
+
async def test_download_info_endpoint_non_200_status(self):
|
|
3387
|
+
"""Test download info endpoint non-200 status with error reading body"""
|
|
3388
|
+
with tempfile.TemporaryDirectory() as tmpdir:
|
|
3389
|
+
with patch("aiohttp.ClientSession") as mock_session:
|
|
3390
|
+
mock_session_instance = AsyncMock()
|
|
3391
|
+
mock_session.return_value.__aenter__ = AsyncMock(
|
|
3392
|
+
return_value=mock_session_instance
|
|
3393
|
+
)
|
|
3394
|
+
mock_session.return_value.__aexit__ = AsyncMock(
|
|
3395
|
+
return_value=None
|
|
3396
|
+
)
|
|
3397
|
+
|
|
3398
|
+
# Mock /info endpoint returning 500 with error reading body
|
|
3399
|
+
async def _get_info(*args, **kwargs):
|
|
3400
|
+
mock_resp = AsyncMock()
|
|
3401
|
+
mock_resp.status = 500
|
|
3402
|
+
mock_resp.text = AsyncMock(
|
|
3403
|
+
side_effect=Exception("Read error")
|
|
3404
|
+
)
|
|
3405
|
+
mock_resp.request_info = Mock()
|
|
3406
|
+
mock_resp.history = ()
|
|
3407
|
+
mock_resp.__aenter__ = AsyncMock(return_value=mock_resp)
|
|
3408
|
+
mock_resp.__aexit__ = AsyncMock(return_value=None)
|
|
3409
|
+
async with mock_resp:
|
|
3410
|
+
if mock_resp.status != 200:
|
|
3411
|
+
try:
|
|
3412
|
+
error_text = await mock_resp.text()
|
|
3413
|
+
except Exception:
|
|
3414
|
+
error_text = f"Unable to read response body (status: {mock_resp.status})"
|
|
3415
|
+
raise ConnectionError(
|
|
3416
|
+
f"Failed to get server info (status {mock_resp.status}): {error_text}"
|
|
3417
|
+
)
|
|
3418
|
+
|
|
3419
|
+
with patch(
|
|
3420
|
+
"proximl.utils.transfer.retry_request",
|
|
3421
|
+
side_effect=_get_info,
|
|
3422
|
+
):
|
|
3423
|
+
with patch(
|
|
3424
|
+
"proximl.utils.transfer.ping_endpoint",
|
|
3425
|
+
new_callable=AsyncMock,
|
|
3426
|
+
):
|
|
3427
|
+
with raises(
|
|
3428
|
+
ConnectionError, match="Failed to get server info"
|
|
3429
|
+
):
|
|
3430
|
+
await specimen.download(
|
|
3431
|
+
"example.com", "token", tmpdir
|
|
3432
|
+
)
|
|
3433
|
+
|
|
3434
|
+
@mark.asyncio
|
|
3435
|
+
async def test_download_non_404_error_in_download(self):
|
|
3436
|
+
"""Test download endpoint non-404 error"""
|
|
3437
|
+
with tempfile.TemporaryDirectory() as tmpdir:
|
|
3438
|
+
with patch("aiohttp.ClientSession") as mock_session:
|
|
3439
|
+
mock_session_instance = AsyncMock()
|
|
3440
|
+
mock_session.return_value.__aenter__ = AsyncMock(
|
|
3441
|
+
return_value=mock_session_instance
|
|
3442
|
+
)
|
|
3443
|
+
mock_session.return_value.__aexit__ = AsyncMock(
|
|
3444
|
+
return_value=None
|
|
3445
|
+
)
|
|
3446
|
+
|
|
3447
|
+
async def mock_get(*args, **kwargs):
|
|
3448
|
+
mock_resp = AsyncMock()
|
|
3449
|
+
mock_resp.status = 200
|
|
3450
|
+
mock_resp.json = AsyncMock(return_value={"archive": False})
|
|
3451
|
+
mock_resp.__aenter__ = AsyncMock(return_value=mock_resp)
|
|
3452
|
+
mock_resp.__aexit__ = AsyncMock(return_value=None)
|
|
3453
|
+
mock_resp.text = AsyncMock(return_value="")
|
|
3454
|
+
return mock_resp
|
|
3455
|
+
|
|
3456
|
+
mock_session_instance.get = AsyncMock(side_effect=mock_get)
|
|
3457
|
+
|
|
3458
|
+
# Mock _download to raise ClientResponseError for non-404 errors
|
|
3459
|
+
async def _download(*args, **kwargs):
|
|
3460
|
+
# Simulate session.get() returning a response with status 500
|
|
3461
|
+
mock_resp = AsyncMock()
|
|
3462
|
+
mock_resp.status = 500
|
|
3463
|
+
mock_resp.text = AsyncMock(return_value="Server error")
|
|
3464
|
+
mock_resp.close = Mock()
|
|
3465
|
+
mock_resp.request_info = Mock()
|
|
3466
|
+
mock_resp.history = ()
|
|
3467
|
+
# This simulates the code in _download that checks status
|
|
3468
|
+
if mock_resp.status != 200:
|
|
3469
|
+
text = await mock_resp.text()
|
|
3470
|
+
mock_resp.close()
|
|
3471
|
+
if mock_resp.status == 404:
|
|
3472
|
+
raise ConnectionError(
|
|
3473
|
+
"Download endpoint not available (404)"
|
|
3474
|
+
)
|
|
3475
|
+
raise ClientResponseError(
|
|
3476
|
+
request_info=mock_resp.request_info,
|
|
3477
|
+
history=mock_resp.history,
|
|
3478
|
+
status=mock_resp.status,
|
|
3479
|
+
message=text,
|
|
3480
|
+
)
|
|
3481
|
+
return mock_resp
|
|
3482
|
+
|
|
3483
|
+
# Mock retry_request to raise ClientResponseError for the download call
|
|
3484
|
+
# First call returns info, second call (for download) raises
|
|
3485
|
+
def mock_retry_side_effect(func, *args, **kwargs):
|
|
3486
|
+
# Check if this is the _download call by checking if func is callable
|
|
3487
|
+
# For the info call, we return the dict
|
|
3488
|
+
# For the download call, we raise
|
|
3489
|
+
if callable(func):
|
|
3490
|
+
# This is likely _download - raise the error
|
|
3491
|
+
raise ClientResponseError(
|
|
3492
|
+
request_info=Mock(),
|
|
3493
|
+
history=(),
|
|
3494
|
+
status=500,
|
|
3495
|
+
message="Server error",
|
|
3496
|
+
)
|
|
3497
|
+
# This shouldn't happen, but return the info dict
|
|
3498
|
+
return {"archive": False}
|
|
3499
|
+
|
|
3500
|
+
with patch(
|
|
3501
|
+
"proximl.utils.transfer.retry_request",
|
|
3502
|
+
side_effect=[
|
|
3503
|
+
{"archive": False},
|
|
3504
|
+
ClientResponseError(
|
|
3505
|
+
request_info=Mock(),
|
|
3506
|
+
history=(),
|
|
3507
|
+
status=500,
|
|
3508
|
+
message="Server error",
|
|
3509
|
+
),
|
|
3510
|
+
],
|
|
3511
|
+
):
|
|
3512
|
+
with patch(
|
|
3513
|
+
"proximl.utils.transfer.ping_endpoint",
|
|
3514
|
+
new_callable=AsyncMock,
|
|
3515
|
+
):
|
|
3516
|
+
with raises(ClientResponseError):
|
|
3517
|
+
await specimen.download(
|
|
3518
|
+
"example.com", "token", tmpdir
|
|
3519
|
+
)
|
|
3520
|
+
|
|
3521
|
+
@mark.asyncio
|
|
3522
|
+
async def test_download_content_disposition_filename_no_quotes_fallback(
|
|
3523
|
+
self,
|
|
3524
|
+
):
|
|
3525
|
+
"""Test download filename parsing without quotes fallback"""
|
|
3526
|
+
with tempfile.TemporaryDirectory() as tmpdir:
|
|
3527
|
+
with patch("aiohttp.ClientSession") as mock_session:
|
|
3528
|
+
mock_session_instance = AsyncMock()
|
|
3529
|
+
mock_session.return_value.__aenter__ = AsyncMock(
|
|
3530
|
+
return_value=mock_session_instance
|
|
3531
|
+
)
|
|
3532
|
+
mock_session.return_value.__aexit__ = AsyncMock(
|
|
3533
|
+
return_value=None
|
|
3534
|
+
)
|
|
3535
|
+
|
|
3536
|
+
def mock_get(*args, **kwargs):
|
|
3537
|
+
url = args[0] if args else kwargs.get("url", "")
|
|
3538
|
+
# Handle /ping endpoint for ping_endpoint
|
|
3539
|
+
if "/ping" in url:
|
|
3540
|
+
mock_resp = AsyncMock()
|
|
3541
|
+
mock_resp.status = 200
|
|
3542
|
+
mock_resp.__aenter__ = AsyncMock(
|
|
3543
|
+
return_value=mock_resp
|
|
3544
|
+
)
|
|
3545
|
+
mock_resp.__aexit__ = AsyncMock(return_value=None)
|
|
3546
|
+
return mock_resp
|
|
3547
|
+
# Handle /info endpoint
|
|
3548
|
+
elif "/info" in url:
|
|
3549
|
+
mock_resp = AsyncMock()
|
|
3550
|
+
mock_resp.status = 200
|
|
3551
|
+
mock_resp.json = AsyncMock(
|
|
3552
|
+
return_value={"archive": True}
|
|
3553
|
+
)
|
|
3554
|
+
mock_resp.text = AsyncMock(return_value="")
|
|
3555
|
+
mock_resp.__aenter__ = AsyncMock(
|
|
3556
|
+
return_value=mock_resp
|
|
3557
|
+
)
|
|
3558
|
+
mock_resp.__aexit__ = AsyncMock(return_value=None)
|
|
3559
|
+
return mock_resp
|
|
3560
|
+
else:
|
|
3561
|
+
# For /download endpoint, return awaitable that resolves to response
|
|
3562
|
+
async def get_download_response():
|
|
3563
|
+
mock_resp = AsyncMock()
|
|
3564
|
+
mock_resp.status = 200
|
|
3565
|
+
mock_resp.headers = {
|
|
3566
|
+
"Content-Type": "application/zip",
|
|
3567
|
+
"Content-Disposition": "attachment; filename=test-file.zip",
|
|
3568
|
+
}
|
|
3569
|
+
|
|
3570
|
+
async def chunk_iter():
|
|
3571
|
+
yield b"zip data"
|
|
3572
|
+
yield b""
|
|
3573
|
+
|
|
3574
|
+
mock_resp.content.iter_chunked = (
|
|
3575
|
+
lambda size: chunk_iter()
|
|
3576
|
+
)
|
|
3577
|
+
mock_resp.close = Mock()
|
|
3578
|
+
return mock_resp
|
|
3579
|
+
|
|
3580
|
+
return get_download_response()
|
|
3581
|
+
|
|
3582
|
+
mock_session_instance.get = Mock(side_effect=mock_get)
|
|
3583
|
+
|
|
3584
|
+
# Mock retry_request to actually call the function passed to it
|
|
3585
|
+
async def mock_retry(func, *args, **kwargs):
|
|
3586
|
+
return await func(*args, **kwargs)
|
|
3587
|
+
|
|
3588
|
+
with patch(
|
|
3589
|
+
"proximl.utils.transfer.retry_request",
|
|
3590
|
+
side_effect=mock_retry,
|
|
3591
|
+
):
|
|
3592
|
+
with patch(
|
|
3593
|
+
"proximl.utils.transfer.ping_endpoint",
|
|
3594
|
+
new_callable=AsyncMock,
|
|
3595
|
+
):
|
|
3596
|
+
mock_file_context = AsyncMock()
|
|
3597
|
+
mock_file_context.__aenter__ = AsyncMock(
|
|
3598
|
+
return_value=AsyncMock()
|
|
3599
|
+
)
|
|
3600
|
+
mock_file_context.__aexit__ = AsyncMock(
|
|
3601
|
+
return_value=None
|
|
3602
|
+
)
|
|
3603
|
+
mock_file_context.__aenter__.return_value.write = (
|
|
3604
|
+
AsyncMock()
|
|
3605
|
+
)
|
|
3606
|
+
with patch(
|
|
3607
|
+
"aiofiles.open", return_value=mock_file_context
|
|
3608
|
+
) as mock_file:
|
|
3609
|
+
mock_finalize_response = AsyncMock()
|
|
3610
|
+
mock_finalize_response.status = 200
|
|
3611
|
+
mock_finalize_response.json = AsyncMock(
|
|
3612
|
+
return_value={"status": "ok"}
|
|
3613
|
+
)
|
|
3614
|
+
mock_finalize_response.__aenter__ = AsyncMock(
|
|
3615
|
+
return_value=mock_finalize_response
|
|
3616
|
+
)
|
|
3617
|
+
mock_finalize_response.__aexit__ = AsyncMock(
|
|
3618
|
+
return_value=None
|
|
3619
|
+
)
|
|
3620
|
+
|
|
3621
|
+
# session.post() should return something that is both awaitable and an async context manager
|
|
3622
|
+
class AwaitableContextManager:
|
|
3623
|
+
def __init__(self, return_value):
|
|
3624
|
+
self.return_value = return_value
|
|
3625
|
+
|
|
3626
|
+
def __await__(self):
|
|
3627
|
+
yield
|
|
3628
|
+
return self
|
|
3629
|
+
|
|
3630
|
+
async def __aenter__(self):
|
|
3631
|
+
return self.return_value
|
|
3632
|
+
|
|
3633
|
+
async def __aexit__(self, *args):
|
|
3634
|
+
return None
|
|
3635
|
+
|
|
3636
|
+
mock_post_context = AwaitableContextManager(
|
|
3637
|
+
mock_finalize_response
|
|
3638
|
+
)
|
|
3639
|
+
mock_session_instance.post = Mock(
|
|
3640
|
+
return_value=mock_post_context
|
|
3641
|
+
)
|
|
3642
|
+
|
|
3643
|
+
with patch(
|
|
3644
|
+
"proximl.utils.transfer.ping_endpoint",
|
|
3645
|
+
new_callable=AsyncMock,
|
|
3646
|
+
):
|
|
3647
|
+
await specimen.download(
|
|
3648
|
+
"example.com", "token", tmpdir
|
|
3649
|
+
)
|
|
3650
|
+
|
|
3651
|
+
@mark.asyncio
|
|
3652
|
+
async def test_download_finalize_success_logging(self):
|
|
3653
|
+
"""Test download finalize success logging"""
|
|
3654
|
+
with tempfile.TemporaryDirectory() as tmpdir:
|
|
3655
|
+
with patch("aiohttp.ClientSession") as mock_session:
|
|
3656
|
+
mock_session_instance = AsyncMock()
|
|
3657
|
+
mock_session.return_value.__aenter__ = AsyncMock(
|
|
3658
|
+
return_value=mock_session_instance
|
|
3659
|
+
)
|
|
3660
|
+
mock_session.return_value.__aexit__ = AsyncMock(
|
|
3661
|
+
return_value=None
|
|
3662
|
+
)
|
|
3663
|
+
|
|
3664
|
+
# Mock /info endpoint returning success
|
|
3665
|
+
mock_info_response = AsyncMock()
|
|
3666
|
+
mock_info_response.status = 200
|
|
3667
|
+
mock_info_response.json = AsyncMock(
|
|
3668
|
+
return_value={"archive": False}
|
|
3669
|
+
)
|
|
3670
|
+
mock_info_response.__aenter__ = AsyncMock(
|
|
3671
|
+
return_value=mock_info_response
|
|
3672
|
+
)
|
|
3673
|
+
mock_info_response.__aexit__ = AsyncMock(return_value=None)
|
|
3674
|
+
|
|
3675
|
+
mock_response = AsyncMock()
|
|
3676
|
+
mock_response.status = 200
|
|
3677
|
+
mock_response.headers = {"Content-Type": "application/x-tar"}
|
|
3678
|
+
|
|
3679
|
+
async def chunk_iter():
|
|
3680
|
+
yield b"tar data"
|
|
3681
|
+
yield b""
|
|
3682
|
+
|
|
3683
|
+
mock_response.content.iter_chunked = lambda size: chunk_iter()
|
|
3684
|
+
mock_response.close = Mock()
|
|
3685
|
+
mock_response.request_info = Mock()
|
|
3686
|
+
mock_response.history = ()
|
|
3687
|
+
|
|
3688
|
+
call_count = 0
|
|
3689
|
+
|
|
3690
|
+
def mock_get(*args, **kwargs):
|
|
3691
|
+
nonlocal call_count
|
|
3692
|
+
url = args[0] if args else kwargs.get("url", "")
|
|
3693
|
+
# Handle /ping endpoint for ping_endpoint
|
|
3694
|
+
if "/ping" in url:
|
|
3695
|
+
mock_resp = AsyncMock()
|
|
3696
|
+
mock_resp.status = 200
|
|
3697
|
+
mock_resp.__aenter__ = AsyncMock(
|
|
3698
|
+
return_value=mock_resp
|
|
3699
|
+
)
|
|
3700
|
+
mock_resp.__aexit__ = AsyncMock(return_value=None)
|
|
3701
|
+
return mock_resp
|
|
3702
|
+
call_count += 1
|
|
3703
|
+
if call_count == 1:
|
|
3704
|
+
# For /info endpoint, return async context manager (used with async with)
|
|
3705
|
+
mock_get_response = AsyncMock()
|
|
3706
|
+
mock_get_response.__aenter__ = AsyncMock(
|
|
3707
|
+
return_value=mock_info_response
|
|
3708
|
+
)
|
|
3709
|
+
mock_get_response.__aexit__ = AsyncMock(
|
|
3710
|
+
return_value=None
|
|
3711
|
+
)
|
|
3712
|
+
return mock_get_response
|
|
3713
|
+
else:
|
|
3714
|
+
# For /download endpoint, return awaitable that resolves to response (used with await)
|
|
3715
|
+
async def get_download_response():
|
|
3716
|
+
return mock_response
|
|
3717
|
+
|
|
3718
|
+
return get_download_response()
|
|
3719
|
+
|
|
3720
|
+
mock_session_instance.get = Mock(side_effect=mock_get)
|
|
3721
|
+
|
|
3722
|
+
# retry_request is called 3 times: _get_info, _download, _finalize
|
|
3723
|
+
async def mock_retry(func, *args, **kwargs):
|
|
3724
|
+
return await func(*args, **kwargs)
|
|
3725
|
+
|
|
3726
|
+
with patch(
|
|
3727
|
+
"proximl.utils.transfer.retry_request",
|
|
3728
|
+
side_effect=mock_retry,
|
|
3729
|
+
):
|
|
3730
|
+
with patch(
|
|
3731
|
+
"proximl.utils.transfer.ping_endpoint",
|
|
3732
|
+
new_callable=AsyncMock,
|
|
3733
|
+
):
|
|
3734
|
+
with patch(
|
|
3735
|
+
"asyncio.create_subprocess_exec"
|
|
3736
|
+
) as mock_subprocess:
|
|
3737
|
+
mock_process = AsyncMock()
|
|
3738
|
+
mock_process.stdin = Mock()
|
|
3739
|
+
mock_process.stdin.write = Mock()
|
|
3740
|
+
mock_process.stdin.drain = AsyncMock()
|
|
3741
|
+
mock_process.stdin.close = Mock()
|
|
3742
|
+
mock_process.returncode = 0
|
|
3743
|
+
mock_process.wait = AsyncMock(return_value=0)
|
|
3744
|
+
mock_process.stderr.read = AsyncMock(
|
|
3745
|
+
return_value=b""
|
|
3746
|
+
)
|
|
3747
|
+
mock_subprocess.return_value = mock_process
|
|
3748
|
+
|
|
3749
|
+
mock_finalize_response = AsyncMock()
|
|
3750
|
+
mock_finalize_response.status = 200
|
|
3751
|
+
mock_finalize_response.json = AsyncMock(
|
|
3752
|
+
return_value={"status": "ok", "files": 10}
|
|
3753
|
+
)
|
|
3754
|
+
mock_finalize_response.__aenter__ = AsyncMock(
|
|
3755
|
+
return_value=mock_finalize_response
|
|
3756
|
+
)
|
|
3757
|
+
mock_finalize_response.__aexit__ = AsyncMock(
|
|
3758
|
+
return_value=None
|
|
3759
|
+
)
|
|
3760
|
+
|
|
3761
|
+
# session.post() should return something that is both awaitable and an async context manager
|
|
3762
|
+
class AwaitableContextManager:
|
|
3763
|
+
def __init__(self, return_value):
|
|
3764
|
+
self.return_value = return_value
|
|
3765
|
+
|
|
3766
|
+
def __await__(self):
|
|
3767
|
+
yield
|
|
3768
|
+
return self
|
|
3769
|
+
|
|
3770
|
+
async def __aenter__(self):
|
|
3771
|
+
return self.return_value
|
|
3772
|
+
|
|
3773
|
+
async def __aexit__(self, *args):
|
|
3774
|
+
return None
|
|
3775
|
+
|
|
3776
|
+
mock_post_context = AwaitableContextManager(
|
|
3777
|
+
mock_finalize_response
|
|
3778
|
+
)
|
|
3779
|
+
mock_session_instance.post = Mock(
|
|
3780
|
+
return_value=mock_post_context
|
|
3781
|
+
)
|
|
3782
|
+
|
|
3783
|
+
with patch("logging.debug") as mock_log:
|
|
3784
|
+
await specimen.download(
|
|
3785
|
+
"example.com", "token", tmpdir
|
|
3786
|
+
)
|
|
3787
|
+
# Verify logging.debug was called for finalize
|
|
3788
|
+
mock_log.assert_called()
|
|
3789
|
+
|
|
3790
|
+
@mark.asyncio
|
|
3791
|
+
async def test_download_info_endpoint_error_direct(self):
|
|
3792
|
+
"""Test download info endpoint error handling (lines 246-259) - direct execution"""
|
|
3793
|
+
with tempfile.TemporaryDirectory() as tmpdir:
|
|
3794
|
+
with patch("aiohttp.ClientSession") as mock_session:
|
|
3795
|
+
mock_session_instance = AsyncMock()
|
|
3796
|
+
mock_session.return_value.__aenter__ = AsyncMock(
|
|
3797
|
+
return_value=mock_session_instance
|
|
3798
|
+
)
|
|
3799
|
+
mock_session.return_value.__aexit__ = AsyncMock(
|
|
3800
|
+
return_value=None
|
|
3801
|
+
)
|
|
3802
|
+
|
|
3803
|
+
# Mock /info endpoint returning 500
|
|
3804
|
+
mock_info_response = AsyncMock()
|
|
3805
|
+
mock_info_response.status = 500
|
|
3806
|
+
mock_info_response.text = AsyncMock(
|
|
3807
|
+
return_value="Server error"
|
|
3808
|
+
)
|
|
3809
|
+
mock_info_response.request_info = Mock()
|
|
3810
|
+
mock_info_response.history = ()
|
|
3811
|
+
mock_info_response.__aenter__ = AsyncMock(
|
|
3812
|
+
return_value=mock_info_response
|
|
3813
|
+
)
|
|
3814
|
+
mock_info_response.__aexit__ = AsyncMock(return_value=None)
|
|
3815
|
+
|
|
3816
|
+
# Make session.get return an async context manager
|
|
3817
|
+
mock_get_response = AsyncMock()
|
|
3818
|
+
mock_get_response.__aenter__ = AsyncMock(
|
|
3819
|
+
return_value=mock_info_response
|
|
3820
|
+
)
|
|
3821
|
+
mock_get_response.__aexit__ = AsyncMock(return_value=None)
|
|
3822
|
+
mock_session_instance.get = Mock(
|
|
3823
|
+
return_value=mock_get_response
|
|
3824
|
+
)
|
|
3825
|
+
|
|
3826
|
+
# Mock retry_request to actually call the function passed to it
|
|
3827
|
+
async def mock_retry(func, *args, **kwargs):
|
|
3828
|
+
return await func(*args, **kwargs)
|
|
3829
|
+
|
|
3830
|
+
with patch(
|
|
3831
|
+
"proximl.utils.transfer.retry_request",
|
|
3832
|
+
side_effect=mock_retry,
|
|
3833
|
+
):
|
|
3834
|
+
with patch(
|
|
3835
|
+
"proximl.utils.transfer.ping_endpoint",
|
|
3836
|
+
new_callable=AsyncMock,
|
|
3837
|
+
):
|
|
3838
|
+
with raises(
|
|
3839
|
+
ConnectionError, match="Failed to get server info"
|
|
3840
|
+
):
|
|
3841
|
+
await specimen.download(
|
|
3842
|
+
"example.com", "token", tmpdir
|
|
3843
|
+
)
|
|
3844
|
+
|
|
3845
|
+
@mark.asyncio
|
|
3846
|
+
async def test_download_info_endpoint_error_reading_body_direct(self):
|
|
3847
|
+
"""Test download info endpoint error reading body (lines 252-255) - direct execution"""
|
|
3848
|
+
with tempfile.TemporaryDirectory() as tmpdir:
|
|
3849
|
+
with patch("aiohttp.ClientSession") as mock_session:
|
|
3850
|
+
mock_session_instance = AsyncMock()
|
|
3851
|
+
mock_session.return_value.__aenter__ = AsyncMock(
|
|
3852
|
+
return_value=mock_session_instance
|
|
3853
|
+
)
|
|
3854
|
+
mock_session.return_value.__aexit__ = AsyncMock(
|
|
3855
|
+
return_value=None
|
|
3856
|
+
)
|
|
3857
|
+
|
|
3858
|
+
# Mock /info endpoint returning 500 with error reading body
|
|
3859
|
+
mock_info_response = AsyncMock()
|
|
3860
|
+
mock_info_response.status = 500
|
|
3861
|
+
mock_info_response.text = AsyncMock(
|
|
3862
|
+
side_effect=Exception("Read error")
|
|
3863
|
+
)
|
|
3864
|
+
mock_info_response.request_info = Mock()
|
|
3865
|
+
mock_info_response.history = ()
|
|
3866
|
+
mock_info_response.__aenter__ = AsyncMock(
|
|
3867
|
+
return_value=mock_info_response
|
|
3868
|
+
)
|
|
3869
|
+
mock_info_response.__aexit__ = AsyncMock(return_value=None)
|
|
3870
|
+
|
|
3871
|
+
# Make session.get return an async context manager
|
|
3872
|
+
def mock_get(*args, **kwargs):
|
|
3873
|
+
url = args[0] if args else kwargs.get("url", "")
|
|
3874
|
+
# Handle /ping endpoint for ping_endpoint
|
|
3875
|
+
if "/ping" in url:
|
|
3876
|
+
mock_resp = AsyncMock()
|
|
3877
|
+
mock_resp.status = 200
|
|
3878
|
+
mock_resp.__aenter__ = AsyncMock(
|
|
3879
|
+
return_value=mock_resp
|
|
3880
|
+
)
|
|
3881
|
+
mock_resp.__aexit__ = AsyncMock(return_value=None)
|
|
3882
|
+
return mock_resp
|
|
3883
|
+
# Handle /info endpoint
|
|
3884
|
+
mock_get_response = AsyncMock()
|
|
3885
|
+
mock_get_response.__aenter__ = AsyncMock(
|
|
3886
|
+
return_value=mock_info_response
|
|
3887
|
+
)
|
|
3888
|
+
mock_get_response.__aexit__ = AsyncMock(return_value=None)
|
|
3889
|
+
return mock_get_response
|
|
3890
|
+
|
|
3891
|
+
mock_session_instance.get = Mock(side_effect=mock_get)
|
|
3892
|
+
|
|
3893
|
+
# Mock retry_request to actually call the function passed to it
|
|
3894
|
+
async def mock_retry(func, *args, **kwargs):
|
|
3895
|
+
return await func(*args, **kwargs)
|
|
3896
|
+
|
|
3897
|
+
with patch(
|
|
3898
|
+
"proximl.utils.transfer.retry_request",
|
|
3899
|
+
side_effect=mock_retry,
|
|
3900
|
+
):
|
|
3901
|
+
with patch(
|
|
3902
|
+
"proximl.utils.transfer.ping_endpoint",
|
|
3903
|
+
new_callable=AsyncMock,
|
|
3904
|
+
):
|
|
3905
|
+
with raises(
|
|
3906
|
+
ConnectionError,
|
|
3907
|
+
match="Failed to get server info.*Unable to read response body",
|
|
3908
|
+
):
|
|
3909
|
+
await specimen.download(
|
|
3910
|
+
"example.com", "token", tmpdir
|
|
3911
|
+
)
|
|
3912
|
+
|
|
3913
|
+
@mark.asyncio
|
|
3914
|
+
async def test_download_endpoint_404_error_direct(self):
|
|
3915
|
+
"""Test download endpoint 404 error (lines 290-298) - direct execution"""
|
|
3916
|
+
with tempfile.TemporaryDirectory() as tmpdir:
|
|
3917
|
+
with patch("aiohttp.ClientSession") as mock_session:
|
|
3918
|
+
mock_session_instance = AsyncMock()
|
|
3919
|
+
mock_session.return_value.__aenter__ = AsyncMock(
|
|
3920
|
+
return_value=mock_session_instance
|
|
3921
|
+
)
|
|
3922
|
+
mock_session.return_value.__aexit__ = AsyncMock(
|
|
3923
|
+
return_value=None
|
|
3924
|
+
)
|
|
3925
|
+
|
|
3926
|
+
# Mock /info endpoint returning success
|
|
3927
|
+
mock_info_response = AsyncMock()
|
|
3928
|
+
mock_info_response.status = 200
|
|
3929
|
+
mock_info_response.json = AsyncMock(
|
|
3930
|
+
return_value={"archive": False}
|
|
3931
|
+
)
|
|
3932
|
+
mock_info_response.__aenter__ = AsyncMock(
|
|
3933
|
+
return_value=mock_info_response
|
|
3934
|
+
)
|
|
3935
|
+
mock_info_response.__aexit__ = AsyncMock(return_value=None)
|
|
3936
|
+
|
|
3937
|
+
# Mock /download endpoint returning 404
|
|
3938
|
+
mock_download_response = AsyncMock()
|
|
3939
|
+
mock_download_response.status = 404
|
|
3940
|
+
mock_download_response.text = AsyncMock(
|
|
3941
|
+
return_value="Not Found"
|
|
3942
|
+
)
|
|
3943
|
+
mock_download_response.close = Mock()
|
|
3944
|
+
mock_download_response.request_info = Mock()
|
|
3945
|
+
mock_download_response.history = ()
|
|
3946
|
+
|
|
3947
|
+
call_count = 0
|
|
3948
|
+
|
|
3949
|
+
def mock_get(*args, **kwargs):
|
|
3950
|
+
nonlocal call_count
|
|
3951
|
+
url = args[0] if args else kwargs.get("url", "")
|
|
3952
|
+
# Handle /ping endpoint for ping_endpoint
|
|
3953
|
+
if "/ping" in url:
|
|
3954
|
+
mock_resp = AsyncMock()
|
|
3955
|
+
mock_resp.status = 200
|
|
3956
|
+
mock_resp.__aenter__ = AsyncMock(
|
|
3957
|
+
return_value=mock_resp
|
|
3958
|
+
)
|
|
3959
|
+
mock_resp.__aexit__ = AsyncMock(return_value=None)
|
|
3960
|
+
return mock_resp
|
|
3961
|
+
call_count += 1
|
|
3962
|
+
if call_count == 1:
|
|
3963
|
+
# For /info endpoint, return async context manager (used with async with)
|
|
3964
|
+
mock_get_response = AsyncMock()
|
|
3965
|
+
mock_get_response.__aenter__ = AsyncMock(
|
|
3966
|
+
return_value=mock_info_response
|
|
3967
|
+
)
|
|
3968
|
+
mock_get_response.__aexit__ = AsyncMock(
|
|
3969
|
+
return_value=None
|
|
3970
|
+
)
|
|
3971
|
+
return mock_get_response
|
|
3972
|
+
else:
|
|
3973
|
+
# For /download endpoint, return awaitable that resolves to response (used with await)
|
|
3974
|
+
async def get_download_response():
|
|
3975
|
+
return mock_download_response
|
|
3976
|
+
|
|
3977
|
+
return get_download_response()
|
|
3978
|
+
|
|
3979
|
+
mock_session_instance.get = Mock(side_effect=mock_get)
|
|
3980
|
+
|
|
3981
|
+
# Mock retry_request to actually call the function passed to it
|
|
3982
|
+
async def mock_retry(func, *args, **kwargs):
|
|
3983
|
+
return await func(*args, **kwargs)
|
|
3984
|
+
|
|
3985
|
+
with patch(
|
|
3986
|
+
"proximl.utils.transfer.retry_request",
|
|
3987
|
+
side_effect=mock_retry,
|
|
3988
|
+
):
|
|
3989
|
+
with patch(
|
|
3990
|
+
"proximl.utils.transfer.ping_endpoint",
|
|
3991
|
+
new_callable=AsyncMock,
|
|
3992
|
+
):
|
|
3993
|
+
# The 404 error should be raised as ClientResponseError and retried, but eventually
|
|
3994
|
+
# it will raise ConnectionError with the message about endpoint not available
|
|
3995
|
+
with raises((ConnectionError, ClientResponseError)):
|
|
3996
|
+
await specimen.download(
|
|
3997
|
+
"example.com", "token", tmpdir
|
|
3998
|
+
)
|
|
3999
|
+
|
|
4000
|
+
@mark.asyncio
|
|
4001
|
+
async def test_download_endpoint_non_404_error_direct(self):
|
|
4002
|
+
"""Test download endpoint non-404 error (lines 290-304) - direct execution"""
|
|
4003
|
+
with tempfile.TemporaryDirectory() as tmpdir:
|
|
4004
|
+
with patch("aiohttp.ClientSession") as mock_session:
|
|
4005
|
+
mock_session_instance = AsyncMock()
|
|
4006
|
+
mock_session.return_value.__aenter__ = AsyncMock(
|
|
4007
|
+
return_value=mock_session_instance
|
|
4008
|
+
)
|
|
4009
|
+
mock_session.return_value.__aexit__ = AsyncMock(
|
|
4010
|
+
return_value=None
|
|
4011
|
+
)
|
|
4012
|
+
|
|
4013
|
+
# Mock /info endpoint returning success
|
|
4014
|
+
mock_info_response = AsyncMock()
|
|
4015
|
+
mock_info_response.status = 200
|
|
4016
|
+
mock_info_response.json = AsyncMock(
|
|
4017
|
+
return_value={"archive": False}
|
|
4018
|
+
)
|
|
4019
|
+
mock_info_response.__aenter__ = AsyncMock(
|
|
4020
|
+
return_value=mock_info_response
|
|
4021
|
+
)
|
|
4022
|
+
mock_info_response.__aexit__ = AsyncMock(return_value=None)
|
|
4023
|
+
|
|
4024
|
+
# Mock /download endpoint returning 500
|
|
4025
|
+
mock_download_response = AsyncMock()
|
|
4026
|
+
mock_download_response.status = 500
|
|
4027
|
+
mock_download_response.text = AsyncMock(
|
|
4028
|
+
return_value="Server error"
|
|
4029
|
+
)
|
|
4030
|
+
mock_download_response.close = Mock()
|
|
4031
|
+
mock_download_response.request_info = Mock()
|
|
4032
|
+
mock_download_response.history = ()
|
|
4033
|
+
|
|
4034
|
+
call_count = 0
|
|
4035
|
+
|
|
4036
|
+
def mock_get(*args, **kwargs):
|
|
4037
|
+
nonlocal call_count
|
|
4038
|
+
url = args[0] if args else kwargs.get("url", "")
|
|
4039
|
+
# Handle /ping endpoint for ping_endpoint
|
|
4040
|
+
if "/ping" in url:
|
|
4041
|
+
mock_resp = AsyncMock()
|
|
4042
|
+
mock_resp.status = 200
|
|
4043
|
+
mock_resp.__aenter__ = AsyncMock(
|
|
4044
|
+
return_value=mock_resp
|
|
4045
|
+
)
|
|
4046
|
+
mock_resp.__aexit__ = AsyncMock(return_value=None)
|
|
4047
|
+
return mock_resp
|
|
4048
|
+
call_count += 1
|
|
4049
|
+
if call_count == 1:
|
|
4050
|
+
# For /info endpoint, return async context manager (used with async with)
|
|
4051
|
+
mock_get_response = AsyncMock()
|
|
4052
|
+
mock_get_response.__aenter__ = AsyncMock(
|
|
4053
|
+
return_value=mock_info_response
|
|
4054
|
+
)
|
|
4055
|
+
mock_get_response.__aexit__ = AsyncMock(
|
|
4056
|
+
return_value=None
|
|
4057
|
+
)
|
|
4058
|
+
return mock_get_response
|
|
4059
|
+
else:
|
|
4060
|
+
# For /download endpoint, return awaitable that resolves to response (used with await)
|
|
4061
|
+
async def get_download_response():
|
|
4062
|
+
return mock_download_response
|
|
4063
|
+
|
|
4064
|
+
return get_download_response()
|
|
4065
|
+
|
|
4066
|
+
mock_session_instance.get = Mock(side_effect=mock_get)
|
|
4067
|
+
|
|
4068
|
+
# Mock retry_request to actually call the function passed to it
|
|
4069
|
+
async def mock_retry(func, *args, **kwargs):
|
|
4070
|
+
return await func(*args, **kwargs)
|
|
4071
|
+
|
|
4072
|
+
with patch(
|
|
4073
|
+
"proximl.utils.transfer.retry_request",
|
|
4074
|
+
side_effect=mock_retry,
|
|
4075
|
+
):
|
|
4076
|
+
with patch(
|
|
4077
|
+
"proximl.utils.transfer.ping_endpoint",
|
|
4078
|
+
new_callable=AsyncMock,
|
|
4079
|
+
):
|
|
4080
|
+
with raises(ClientResponseError) as exc_info:
|
|
4081
|
+
await specimen.download(
|
|
4082
|
+
"example.com", "token", tmpdir
|
|
4083
|
+
)
|
|
4084
|
+
assert exc_info.value.status == 500
|
|
4085
|
+
|
|
4086
|
+
@mark.asyncio
|
|
4087
|
+
async def test_download_filename_fallback_no_quotes_direct(self):
|
|
4088
|
+
"""Test download filename parsing fallback without quotes (lines 339-343) - direct execution"""
|
|
4089
|
+
with tempfile.TemporaryDirectory() as tmpdir:
|
|
4090
|
+
with patch("aiohttp.ClientSession") as mock_session:
|
|
4091
|
+
mock_session_instance = AsyncMock()
|
|
4092
|
+
mock_session.return_value.__aenter__ = AsyncMock(
|
|
4093
|
+
return_value=mock_session_instance
|
|
4094
|
+
)
|
|
4095
|
+
mock_session.return_value.__aexit__ = AsyncMock(
|
|
4096
|
+
return_value=None
|
|
4097
|
+
)
|
|
4098
|
+
|
|
4099
|
+
# Mock /info endpoint returning success
|
|
4100
|
+
mock_info_response = AsyncMock()
|
|
4101
|
+
mock_info_response.status = 200
|
|
4102
|
+
mock_info_response.json = AsyncMock(
|
|
4103
|
+
return_value={"archive": True}
|
|
4104
|
+
)
|
|
4105
|
+
mock_info_response.__aenter__ = AsyncMock(
|
|
4106
|
+
return_value=mock_info_response
|
|
4107
|
+
)
|
|
4108
|
+
mock_info_response.__aexit__ = AsyncMock(return_value=None)
|
|
4109
|
+
|
|
4110
|
+
# Mock /download endpoint with Content-Disposition that forces fallback regex
|
|
4111
|
+
# To trigger the fallback path (lines 339-343), we need the first regex to not match
|
|
4112
|
+
# The first regex is r'filename="?([^"]+)"?' which matches filename=value
|
|
4113
|
+
# To make it not match, we'll patch the regex search to return None for the first attempt
|
|
4114
|
+
mock_download_response = AsyncMock()
|
|
4115
|
+
mock_download_response.status = 200
|
|
4116
|
+
mock_download_response.headers = {
|
|
4117
|
+
"Content-Type": "application/zip",
|
|
4118
|
+
"Content-Disposition": "attachment; filename=test-file.zip",
|
|
4119
|
+
}
|
|
4120
|
+
|
|
4121
|
+
async def chunk_iter():
|
|
4122
|
+
yield b"zip data"
|
|
4123
|
+
yield b""
|
|
4124
|
+
|
|
4125
|
+
mock_download_response.content.iter_chunked = (
|
|
4126
|
+
lambda size: chunk_iter()
|
|
4127
|
+
)
|
|
4128
|
+
mock_download_response.close = Mock()
|
|
4129
|
+
mock_download_response.request_info = Mock()
|
|
4130
|
+
mock_download_response.history = ()
|
|
4131
|
+
|
|
4132
|
+
call_count = 0
|
|
4133
|
+
|
|
4134
|
+
def mock_get(*args, **kwargs):
|
|
4135
|
+
nonlocal call_count
|
|
4136
|
+
url = args[0] if args else kwargs.get("url", "")
|
|
4137
|
+
# Handle /ping endpoint for ping_endpoint
|
|
4138
|
+
if "/ping" in url:
|
|
4139
|
+
mock_resp = AsyncMock()
|
|
4140
|
+
mock_resp.status = 200
|
|
4141
|
+
mock_resp.__aenter__ = AsyncMock(
|
|
4142
|
+
return_value=mock_resp
|
|
4143
|
+
)
|
|
4144
|
+
mock_resp.__aexit__ = AsyncMock(return_value=None)
|
|
4145
|
+
return mock_resp
|
|
4146
|
+
call_count += 1
|
|
4147
|
+
if call_count == 1:
|
|
4148
|
+
# For /info endpoint, return async context manager (used with async with)
|
|
4149
|
+
mock_get_response = AsyncMock()
|
|
4150
|
+
mock_get_response.__aenter__ = AsyncMock(
|
|
4151
|
+
return_value=mock_info_response
|
|
4152
|
+
)
|
|
4153
|
+
mock_get_response.__aexit__ = AsyncMock(
|
|
4154
|
+
return_value=None
|
|
4155
|
+
)
|
|
4156
|
+
return mock_get_response
|
|
4157
|
+
else:
|
|
4158
|
+
# For /download endpoint, return awaitable that resolves to response (used with await)
|
|
4159
|
+
async def get_download_response():
|
|
4160
|
+
return mock_download_response
|
|
4161
|
+
|
|
4162
|
+
return get_download_response()
|
|
4163
|
+
|
|
4164
|
+
mock_session_instance.get = Mock(side_effect=mock_get)
|
|
4165
|
+
|
|
4166
|
+
# Mock retry_request to actually call the function passed to it
|
|
4167
|
+
async def mock_retry(func, *args, **kwargs):
|
|
4168
|
+
return await func(*args, **kwargs)
|
|
4169
|
+
|
|
4170
|
+
with patch(
|
|
4171
|
+
"proximl.utils.transfer.retry_request",
|
|
4172
|
+
side_effect=mock_retry,
|
|
4173
|
+
):
|
|
4174
|
+
with patch(
|
|
4175
|
+
"proximl.utils.transfer.ping_endpoint",
|
|
4176
|
+
new_callable=AsyncMock,
|
|
4177
|
+
):
|
|
4178
|
+
# Patch re.search to return None for the first regex call (with quotes pattern)
|
|
4179
|
+
# This forces the fallback regex to be used (lines 339-343)
|
|
4180
|
+
original_search = re.search
|
|
4181
|
+
search_call_count = 0
|
|
4182
|
+
|
|
4183
|
+
def mock_re_search(pattern, string, *args, **kwargs):
|
|
4184
|
+
nonlocal search_call_count
|
|
4185
|
+
search_call_count += 1
|
|
4186
|
+
# For the first call (the quotes regex), return None to force fallback
|
|
4187
|
+
if (
|
|
4188
|
+
search_call_count == 1
|
|
4189
|
+
and 'filename="?([^"]+)"?' in pattern
|
|
4190
|
+
):
|
|
4191
|
+
return None
|
|
4192
|
+
# For subsequent calls, use the real re.search
|
|
4193
|
+
return original_search(
|
|
4194
|
+
pattern, string, *args, **kwargs
|
|
4195
|
+
)
|
|
4196
|
+
|
|
4197
|
+
mock_file_context = AsyncMock()
|
|
4198
|
+
mock_file_context.__aenter__ = AsyncMock(
|
|
4199
|
+
return_value=AsyncMock()
|
|
4200
|
+
)
|
|
4201
|
+
mock_file_context.__aexit__ = AsyncMock(
|
|
4202
|
+
return_value=None
|
|
4203
|
+
)
|
|
4204
|
+
mock_file_context.__aenter__.return_value.write = (
|
|
4205
|
+
AsyncMock()
|
|
4206
|
+
)
|
|
4207
|
+
|
|
4208
|
+
mock_finalize_response = AsyncMock()
|
|
4209
|
+
mock_finalize_response.status = 200
|
|
4210
|
+
mock_finalize_response.json = AsyncMock(
|
|
4211
|
+
return_value={"status": "ok"}
|
|
4212
|
+
)
|
|
4213
|
+
mock_finalize_response.__aenter__ = AsyncMock(
|
|
4214
|
+
return_value=mock_finalize_response
|
|
4215
|
+
)
|
|
4216
|
+
mock_finalize_response.__aexit__ = AsyncMock(
|
|
4217
|
+
return_value=None
|
|
4218
|
+
)
|
|
4219
|
+
|
|
4220
|
+
# session.post() should return something that is both awaitable and an async context manager
|
|
4221
|
+
class AwaitableContextManager:
|
|
4222
|
+
def __init__(self, return_value):
|
|
4223
|
+
self.return_value = return_value
|
|
4224
|
+
|
|
4225
|
+
def __await__(self):
|
|
4226
|
+
yield
|
|
4227
|
+
return self
|
|
4228
|
+
|
|
4229
|
+
async def __aenter__(self):
|
|
4230
|
+
return self.return_value
|
|
4231
|
+
|
|
4232
|
+
async def __aexit__(self, *args):
|
|
4233
|
+
return None
|
|
4234
|
+
|
|
4235
|
+
mock_post_context = AwaitableContextManager(
|
|
4236
|
+
mock_finalize_response
|
|
4237
|
+
)
|
|
4238
|
+
mock_session_instance.post = Mock(
|
|
4239
|
+
return_value=mock_post_context
|
|
4240
|
+
)
|
|
4241
|
+
|
|
4242
|
+
with patch(
|
|
4243
|
+
"proximl.utils.transfer.retry_request",
|
|
4244
|
+
side_effect=mock_retry,
|
|
4245
|
+
):
|
|
4246
|
+
with patch(
|
|
4247
|
+
"proximl.utils.transfer.re.search",
|
|
4248
|
+
side_effect=mock_re_search,
|
|
4249
|
+
):
|
|
4250
|
+
with patch(
|
|
4251
|
+
"aiofiles.open",
|
|
4252
|
+
return_value=mock_file_context,
|
|
4253
|
+
):
|
|
4254
|
+
await specimen.download(
|
|
4255
|
+
"example.com", "token", tmpdir
|
|
4256
|
+
)
|
|
4257
|
+
# Verify the file was written (filename parsing worked)
|
|
4258
|
+
assert (
|
|
4259
|
+
mock_file_context.__aenter__.return_value.write.called
|
|
4260
|
+
)
|