Skip to content

Commit 828f810

Browse files
rbell517rbell517
committed
Retry template cleanup and more test robustness
1 parent 75f2e5b commit 828f810

6 files changed

Lines changed: 141 additions & 35 deletions

File tree

‎nisystemlink/clients/core/_uplink/_multipart_retry.py‎

Lines changed: 33 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -95,6 +95,11 @@ def __init__(self) -> None:
9595
int, tuple[type[BaseException], BaseException, Any]
9696
] = {}
9797

98+
def _clear_request_state(self, request_id: int) -> None:
99+
self._attempted_request_ids.discard(request_id)
100+
self._responses_by_request_id.pop(request_id, None)
101+
self._exceptions_by_request_id.pop(request_id, None)
102+
98103
def before_request(self, request: tuple[str, str, dict[str, Any]]) -> Any:
99104
_, _, extras = request
100105
request_id = id(request)
@@ -104,10 +109,12 @@ def before_request(self, request: tuple[str, str, dict[str, Any]]) -> Any:
104109

105110
for part in extras.get("files", {}).values():
106111
if _rewind_retryable_part(part) is _RewindResult.FAILED:
107-
return _get_saved_retry_action(
112+
retry_action = _get_saved_retry_action(
108113
self._responses_by_request_id.get(request_id),
109114
self._exceptions_by_request_id.get(request_id),
110115
)
116+
self._clear_request_state(request_id)
117+
return retry_action
111118
return None
112119

113120
def after_response(
@@ -131,13 +138,36 @@ def after_exception(
131138
return None
132139

133140

141+
class _RetryableMultipartCleanupTemplate(RequestTemplate):
142+
def __init__(self, retry_template: _RetryableMultipartRequestTemplate) -> None:
143+
self._retry_template = retry_template
144+
145+
def after_response(
146+
self, request: tuple[str, str, dict[str, Any]], response: Response
147+
) -> None:
148+
self._retry_template._clear_request_state(id(request))
149+
return None
150+
151+
def after_exception(
152+
self,
153+
request: tuple[str, str, dict[str, Any]],
154+
exc_type: type[BaseException],
155+
exc_val: BaseException,
156+
exc_tb: Any,
157+
) -> None:
158+
self._retry_template._clear_request_state(id(request))
159+
return None
160+
161+
134162
class _RetryableMultipartRequest(decorators.MethodAnnotation):
135163
def modify_request(self, request_builder: Any) -> None:
164+
retryable_template = _RetryableMultipartRequestTemplate()
136165
# Insert ahead of Uplink's retry template so this helper can see the
137166
# original retry-triggering response/exception and short-circuit future
138167
# attempts when a multipart stream cannot be rewound safely.
139-
request_builder._request_templates.insert(
140-
0, _RetryableMultipartRequestTemplate()
168+
request_builder._request_templates.insert(0, retryable_template)
169+
request_builder._request_templates.append(
170+
_RetryableMultipartCleanupTemplate(retryable_template)
141171
)
142172

143173

‎tests/core/test_multipart_retry.py‎

Lines changed: 101 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,12 +6,15 @@
66
from nisystemlink.clients.core import ApiException
77
from nisystemlink.clients.core._uplink._base_client import _handle_http_status
88
from nisystemlink.clients.core._uplink._multipart_retry import (
9+
_RetryableMultipartCleanupTemplate,
910
_RetryableMultipartRequestTemplate,
1011
retryable_multipart_request,
1112
)
1213
from requests import Response
1314
from uplink import Consumer, Part, post, retry
15+
from uplink.clients.io import CompositeRequestTemplate
1416
from uplink.clients.io import state as uplink_state
17+
from uplink.clients.io.interfaces import RequestTemplate
1518

1619

1720
class _NonSeekableStream:
@@ -33,6 +36,18 @@ def seek(self, offset: int) -> None:
3336
raise OSError("cannot rewind")
3437

3538

39+
class _StaticTransitionTemplate(RequestTemplate):
40+
def __init__(self, response_transition=None, exception_transition=None) -> None:
41+
self._response_transition = response_transition
42+
self._exception_transition = exception_transition
43+
44+
def after_response(self, request, response):
45+
return self._response_transition
46+
47+
def after_exception(self, request, exc_type, exc_val, exc_tb):
48+
return self._exception_transition
49+
50+
3651
@retry(
3752
when=retry.when.status(429),
3853
stop=retry.stop.after_attempt(2),
@@ -139,6 +154,92 @@ def test__before_request_when_rewind_fails_after_response__finishes_with_saved_r
139154
next_state = action(uplink_state.BeforeRequest(request))
140155
assert isinstance(next_state, uplink_state.Finish)
141156
assert next_state.response is response
157+
request_id = id(request)
158+
assert request_id not in template._attempted_request_ids
159+
assert request_id not in template._responses_by_request_id
160+
assert request_id not in template._exceptions_by_request_id
161+
162+
def test__terminal_response_after_retry_pipeline__clears_saved_retry_state(self):
163+
request = (
164+
"POST",
165+
"https://example.com/upload",
166+
{
167+
"files": {
168+
"artifact": ("artifact.bin", io.BytesIO(b"artifact")),
169+
}
170+
},
171+
)
172+
response = Response()
173+
response.status_code = 200
174+
response.url = "https://example.com/upload"
175+
template = _RetryableMultipartRequestTemplate()
176+
composite = CompositeRequestTemplate(
177+
[template, _StaticTransitionTemplate(), _RetryableMultipartCleanupTemplate(template)]
178+
)
179+
180+
template.before_request(request)
181+
composite.after_response(request, response)
182+
183+
request_id = id(request)
184+
assert request_id not in template._attempted_request_ids
185+
assert request_id not in template._responses_by_request_id
186+
assert request_id not in template._exceptions_by_request_id
187+
188+
def test__retry_transition_after_response__preserves_saved_retry_state(self):
189+
request = (
190+
"POST",
191+
"https://example.com/upload",
192+
{
193+
"files": {
194+
"artifact": ("artifact.bin", io.BytesIO(b"artifact")),
195+
}
196+
},
197+
)
198+
response = Response()
199+
response.status_code = 503
200+
response.url = "https://example.com/upload"
201+
retry_transition = object()
202+
template = _RetryableMultipartRequestTemplate()
203+
composite = CompositeRequestTemplate(
204+
[
205+
template,
206+
_StaticTransitionTemplate(response_transition=retry_transition),
207+
_RetryableMultipartCleanupTemplate(template),
208+
]
209+
)
210+
211+
template.before_request(request)
212+
213+
assert composite.after_response(request, response) is retry_transition
214+
215+
request_id = id(request)
216+
assert request_id in template._attempted_request_ids
217+
assert template._responses_by_request_id[request_id] is response
218+
assert request_id not in template._exceptions_by_request_id
219+
220+
def test__terminal_exception_after_retry_pipeline__clears_saved_retry_state(self):
221+
request = (
222+
"POST",
223+
"https://example.com/upload",
224+
{
225+
"files": {
226+
"artifact": ("artifact.bin", io.BytesIO(b"artifact")),
227+
}
228+
},
229+
)
230+
exception = RuntimeError("boom")
231+
template = _RetryableMultipartRequestTemplate()
232+
composite = CompositeRequestTemplate(
233+
[template, _StaticTransitionTemplate(), _RetryableMultipartCleanupTemplate(template)]
234+
)
235+
236+
template.before_request(request)
237+
composite.after_exception(request, RuntimeError, exception, None)
238+
239+
request_id = id(request)
240+
assert request_id not in template._attempted_request_ids
241+
assert request_id not in template._responses_by_request_id
242+
assert request_id not in template._exceptions_by_request_id
142243

143244
def test__before_request_on_retry_with_string_only_parts__allows_retry(self):
144245
request = (

‎tests/integration/artifact/test_artifact.py‎

Lines changed: 1 addition & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,6 @@
88
UploadArtifactResponse,
99
)
1010
from nisystemlink.clients.core._http_configuration import HttpConfiguration
11-
from responses import PassthroughResponse
1211
from responses.registries import OrderedRegistry
1312
from uplink.clients.io import blocking_strategy as uplink_blocking_strategy
1413

@@ -69,12 +68,7 @@ def test__upload_artifact_after_rate_limit_retry__artifact_uploaded(
6968
f"{BASE_URL}/ninbartifact/v1/artifacts",
7069
status=429,
7170
)
72-
request_mock.add(
73-
PassthroughResponse(
74-
responses.POST,
75-
f"{BASE_URL}/ninbartifact/v1/artifacts",
76-
)
77-
)
71+
request_mock.add_passthru(f"{BASE_URL}/ninbartifact/v1/artifacts")
7872

7973
upload_response: UploadArtifactResponse = create_artifact()
8074

‎tests/integration/feeds/test_feeds_client.py‎

Lines changed: 2 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,6 @@
1010
from nisystemlink.clients.core import ApiException
1111
from nisystemlink.clients.feeds import FeedsClient
1212
from nisystemlink.clients.feeds.models import CreateFeedRequest, Platform
13-
from responses import PassthroughResponse
1413
from responses.registries import OrderedRegistry
1514
from uplink.clients.io import blocking_strategy as uplink_blocking_strategy
1615

@@ -267,11 +266,8 @@ def test__upload_package_content_after_rate_limit_retry__upload_package_content_
267266
f"{BASE_URL}/nifeed/v1/feeds/{create_feed_resp.id}/packages",
268267
status=429,
269268
)
270-
request_mock.add(
271-
PassthroughResponse(
272-
responses.POST,
273-
f"{BASE_URL}/nifeed/v1/feeds/{create_feed_resp.id}/packages",
274-
)
269+
request_mock.add_passthru(
270+
f"{BASE_URL}/nifeed/v1/feeds/{create_feed_resp.id}/packages"
275271
)
276272

277273
with open(PACKAGE_PATH, "rb") as package:

‎tests/integration/file/test_file_client.py‎

Lines changed: 2 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,6 @@
2121
UpdateMetadataRequest,
2222
)
2323
from nisystemlink.clients.file.utilities import rename_file
24-
from responses import PassthroughResponse
2524
from responses.registries import OrderedRegistry
2625
from uplink.clients.io import blocking_strategy as uplink_blocking_strategy
2726

@@ -137,11 +136,8 @@ def test__upload_file_after_rate_limit_retry__upload_file_succeeds(
137136
f"{BASE_URL}/nifile/v1/service-groups/Default/upload-files",
138137
status=429,
139138
)
140-
request_mock.add(
141-
PassthroughResponse(
142-
responses.POST,
143-
f"{BASE_URL}/nifile/v1/service-groups/Default/upload-files",
144-
)
139+
request_mock.add_passthru(
140+
f"{BASE_URL}/nifile/v1/service-groups/Default/upload-files"
145141
)
146142

147143
file_id = client.upload_file(file=test_file)

‎tests/integration/notebook/test_notebook_client.py‎

Lines changed: 2 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,6 @@
1414
QueryExecutionsRequest,
1515
QueryNotebookRequest,
1616
)
17-
from responses import PassthroughResponse
1817
from responses.registries import OrderedRegistry
1918
from uplink.clients.io import blocking_strategy as uplink_blocking_strategy
2019

@@ -305,12 +304,7 @@ def test__create_notebook_after_rate_limit_retry__notebook_created_with_valid_me
305304
f"{BASE_URL}/ninotebook/v1/notebook",
306305
status=429,
307306
)
308-
request_mock.add(
309-
PassthroughResponse(
310-
responses.POST,
311-
f"{BASE_URL}/ninotebook/v1/notebook",
312-
)
313-
)
307+
request_mock.add_passthru(f"{BASE_URL}/ninotebook/v1/notebook")
314308

315309
with open("tests/integration/notebook/sample_file.ipynb", "rb") as file:
316310
notebook = client.create_notebook(metadata=metadata, content=file)
@@ -351,12 +345,7 @@ def test__update_notebook_metadata_after_rate_limit_retry__update_notebook_metad
351345
f"{BASE_URL}/ninotebook/v1/notebook/{notebook.id}",
352346
status=429,
353347
)
354-
request_mock.add(
355-
PassthroughResponse(
356-
responses.PUT,
357-
f"{BASE_URL}/ninotebook/v1/notebook/{notebook.id}",
358-
)
359-
)
348+
request_mock.add_passthru(f"{BASE_URL}/ninotebook/v1/notebook/{notebook.id}")
360349

361350
response = client.update_notebook(id=notebook.id, metadata=notebook)
362351

0 commit comments

Comments
 (0)