Prove Python TUS retry state transitions

kvz · kvz · commit 6d8c8f179b0e · 2026-06-10T01:51:41.000+02:00
diff --git a/examples/api2-devdock-tus-retry-state-transitions/main.py b/examples/api2-devdock-tus-retry-state-transitions/main.py
@@ -0,0 +1,152 @@
+"""Prove TUS retry attempt state resets after recovered progress."""
+
+import sys
+from io import BytesIO
+from pathlib import Path
+
+from tusclient import client as tus
+
+sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
+from api2devdock import (
+    TusConformancePlanServer,
+    conformance_input_options,
+    conformance_input_source_bytes,
+    conformance_retry_decisions,
+    fail,
+    int_array_value,
+    load_scenario,
+    object_value,
+    scenario_id,
+    string_value,
+    write_result,
+)
+
+
+class RetryStateObserver:
+    def __init__(self, retry_decisions, retry_delays):
+        self.retry_decisions = retry_decisions
+        self.retry_delays = retry_delays
+        self.events = []
+        self.index = 0
+
+    def on_should_retry(self, error, retry_attempt):
+        if self.index >= len(self.retry_decisions):
+            fail(
+                "retry state scenario observed unexpected retry attempt {}".format(
+                    retry_attempt,
+                )
+            )
+
+        expected = self.retry_decisions[self.index]
+        if retry_attempt != expected["retryAttempt"]:
+            fail(
+                "retry state scenario expected retry attempt {}, got {}".format(
+                    expected["retryAttempt"],
+                    retry_attempt,
+                )
+            )
+
+        decision = expected["decision"]
+        self.events.append(
+            {
+                "decision": decision,
+                "kind": "should-retry",
+                "retryAttempt": retry_attempt,
+            }
+        )
+
+        if decision:
+            if retry_attempt >= len(self.retry_delays):
+                fail(
+                    "retry state scenario has no delay for retry attempt {}".format(
+                        retry_attempt,
+                    )
+                )
+            self.events.append(
+                {
+                    "delay": self.retry_delays[retry_attempt],
+                    "kind": "retry-schedule",
+                }
+            )
+
+        self.index += 1
+        return decision
+
+    def assert_complete(self):
+        if self.index == len(self.retry_decisions):
+            return
+
+        fail(
+            "retry state scenario observed {} retry decision(s), expected {}".format(
+                self.index,
+                len(self.retry_decisions),
+            )
+        )
+
+
+def upload_with_retry_state_transitions(conformance_scenario):
+    input_options = conformance_input_options(conformance_scenario)
+    content = conformance_input_source_bytes(conformance_scenario)
+    endpoint_url = string_value(input_options["endpointUrl"], "endpointUrl")
+    metadata = object_value(input_options["metadata"], "metadata")
+    retry_delays = int_array_value(input_options["retryDelays"], "retryDelays")
+    retry_decisions = conformance_retry_decisions(conformance_scenario)
+    completion = object_value(
+        conformance_scenario["completion"],
+        "conformanceScenario.completion",
+    )
+    completion_kind = string_value(
+        completion["kind"],
+        "conformanceScenario.completion.kind",
+    )
+    observer = RetryStateObserver(retry_decisions, retry_delays)
+
+    with TusConformancePlanServer(conformance_scenario, endpoint_url) as conformance_server:
+        client = tus.TusClient(conformance_server.endpoint_url())
+        uploader = client.uploader(
+            file_stream=BytesIO(content),
+            metadata=metadata,
+            on_should_retry=observer.on_should_retry,
+            retry_delays=retry_delays,
+        )
+        uploader.upload()
+        observer.assert_complete()
+
+        if not uploader.url:
+            fail("retry state scenario did not expose an upload URL")
+        if uploader.offset != len(content):
+            fail(
+                "retry state scenario upload offset {}, expected {}".format(
+                    uploader.offset,
+                    len(content),
+                )
+            )
+
+        conformance_server.assert_exhausted()
+        result = conformance_server.result()
+        result["completionKind"] = completion_kind
+        result["errorCalled"] = False
+        result["eventCount"] = len(observer.events)
+        result["events"] = observer.events
+        result["successCalled"] = True
+        result["uploadUrl"] = conformance_server.canonical_url(uploader.url)
+        return result
+
+
+def main():
+    scenario = load_scenario(Path(__file__).with_name("api2-scenario.json"))
+    conformance_scenario = object_value(
+        scenario["conformanceScenario"],
+        "conformanceScenario",
+    )
+    result = upload_with_retry_state_transitions(conformance_scenario)
+    write_result(result)
+    print(
+        "Python TUS SDK devdock scenario {} proved retry state transitions".format(
+            scenario_id(scenario),
+        )
+    )
+
+
+if __name__ == "__main__":
+    main()
diff --git a/examples/api2devdock.py b/examples/api2devdock.py
@@ -139,6 +139,33 @@ def conformance_input_source_kind(conformance_scenario):
     return string_value(input_source["kind"], "conformanceScenario.inputSource.kind")
 
 
+def conformance_retry_decisions(conformance_scenario):
+    decisions = conformance_scenario["retryDecisions"]
+    if not isinstance(decisions, list):
+        fail("conformanceScenario.retryDecisions must be a list")
+
+    result = []
+    for index, decision in enumerate(decisions):
+        decision = object_value(
+            decision,
+            "conformanceScenario.retryDecisions[{}]".format(index),
+        )
+        result.append(
+            {
+                "decision": bool_value(
+                    decision["decision"],
+                    "conformanceScenario.retryDecisions[{}].decision".format(index),
+                ),
+                "retryAttempt": int_value(
+                    decision["retryAttempt"],
+                    "conformanceScenario.retryDecisions[{}].retryAttempt".format(index),
+                ),
+            }
+        )
+
+    return result
+
+
 def conformance_scenario_wants_event(conformance_scenario, event_kind):
     events = conformance_scenario.get("events", [])
     if not isinstance(events, list):
diff --git a/tests/test_uploader.py b/tests/test_uploader.py
@@ -297,6 +297,74 @@ def test_upload_retry(self, request_mock):
             self.uploader.upload_chunk()
         self.assertEqual(self.uploader._retried, num_of_retries)
 
+    @mock.patch('tusclient.uploader.uploader.TusRequest')
+    def test_upload_retry_delays_and_should_retry_reset_after_progress(self, request_mock):
+        first_failure = mock.Mock()
+        first_failure.status_code = 500
+        first_failure.response_content = b''
+        first_failure.response_headers = {}
+        first_failure.perform.return_value = None
+
+        second_failure = mock.Mock()
+        second_failure.status_code = 500
+        second_failure.response_content = b''
+        second_failure.response_headers = {}
+        second_failure.perform.return_value = None
+
+        success = mock.Mock()
+        success.status_code = 204
+        success.response_content = b''
+        success.response_headers = {'upload-offset': '11'}
+        success.perform.return_value = None
+
+        retry_attempts = []
+
+        def should_retry(error, retry_attempt):
+            retry_attempts.append(retry_attempt)
+            return True
+
+        self.uploader.retries = 0
+        self.uploader.retry_delays = [250]
+        self.uploader.on_should_retry = should_retry
+        self.uploader.offset = 0
+        request_mock.side_effect = [first_failure, second_failure, success]
+
+        with mock.patch.object(self.uploader, 'get_offset', side_effect=[5, 5]) as get_offset:
+            with mock.patch('tusclient.uploader.uploader.time.sleep') as sleep:
+                self.uploader.upload_chunk()
+
+        self.assertEqual(retry_attempts, [0, 0])
+        self.assertEqual(self.uploader.offset, 11)
+        self.assertEqual(get_offset.call_count, 2)
+        sleep.assert_has_calls([mock.call(0.25), mock.call(0.25)])
+
+    @mock.patch('tusclient.uploader.uploader.TusRequest')
+    def test_upload_retry_stops_when_should_retry_returns_false(self, request_mock):
+        failure = mock.Mock()
+        failure.status_code = 500
+        failure.response_content = b''
+        failure.response_headers = {}
+        failure.perform.return_value = None
+
+        retry_attempts = []
+
+        def should_retry(error, retry_attempt):
+            retry_attempts.append(retry_attempt)
+            return False
+
+        self.uploader.retries = 0
+        self.uploader.retry_delays = [0]
+        self.uploader.on_should_retry = should_retry
+        request_mock.side_effect = [failure]
+
+        with mock.patch.object(self.uploader, 'get_offset') as get_offset:
+            with pytest.raises(exceptions.TusCommunicationError):
+                self.uploader.upload_chunk()
+
+        self.assertEqual(retry_attempts, [0])
+        self.assertEqual(self.uploader._retried, 0)
+        get_offset.assert_not_called()
+
     @responses.activate
     def test_upload_empty(self):
         responses.add(
diff --git a/tusclient/uploader/baseuploader.py b/tusclient/uploader/baseuploader.py
@@ -87,6 +87,8 @@ class BaseUploader:
             Callback invoked with bytes sent and total bytes after upload progress changes.
         - on_chunk_complete (Optional[Callable]):
             Callback invoked with chunk size, accepted offset, and total bytes after a chunk is accepted.
+        - on_should_retry (Optional[Callable]):
+            Callback invoked with an error and retry attempt before scheduling a retry.
 
     :Constructor Args:
         - file_path (str)
@@ -139,6 +141,7 @@ def __init__(
         parallel_upload_boundaries=None,
         protocol: Optional[str] = None,
         retry_delays=None,
+        on_should_retry: Optional[Callable[[Exception, int], bool]] = None,
         on_progress: Optional[Callable[[int, Optional[int]], None]] = None,
         on_chunk_complete: Optional[Callable[[int, int, Optional[int]], None]] = None,
     ):
@@ -196,6 +199,7 @@ def __init__(
         self.parallel_upload_boundaries = parallel_upload_boundaries
         self.override_patch_method = override_patch_method
         self.retry_delays = retry_delays
+        self.on_should_retry = on_should_retry
         self.on_progress = on_progress
         self.on_chunk_complete = on_chunk_complete
         self._abort_requested = Event()
@@ -232,6 +236,24 @@ def request_method_input_options(self):
             "override_patch_method": self.override_patch_method,
         }
 
+    def _retry_limit(self):
+        if self.retry_delays is not None:
+            return len(self.retry_delays)
+
+        return self.retries
+
+    def _retry_delay_seconds(self, retry_attempt):
+        if self.retry_delays is not None:
+            return self.retry_delays[retry_attempt] / 1000
+
+        return self.retry_delay
+
+    def _should_retry(self, error, retry_attempt):
+        if self.on_should_retry is None:
+            return True
+
+        return bool(self.on_should_retry(error, retry_attempt))
+
     def run_before_request(self, method, url, headers):
         if self.client is not None:
             self.client._set_current_uploader(self)
diff --git a/tusclient/uploader/uploader.py b/tusclient/uploader/uploader.py
@@ -208,19 +208,25 @@ def _do_request(self):
             self._retry_or_cry(error)
 
     def _retry_or_cry(self, error):
-        if self.retries > self._retried:
-            time.sleep(self.retry_delay)
-
-            self._retried += 1
-            try:
-                self.offset = self.get_offset()
-            except TusCommunicationError as err:
-                self._retry_or_cry(err)
-            else:
-                self._do_request()
-        else:
+        retry_attempt = self._retried
+        if self._retry_limit() <= retry_attempt:
+            raise error
+        if not self._should_retry(error, retry_attempt):
             raise error
 
+        time.sleep(self._retry_delay_seconds(retry_attempt))
+        self._retried += 1
+        previous_offset = self.offset
+        try:
+            recovered_offset = self.get_offset()
+        except TusCommunicationError as err:
+            self._retry_or_cry(err)
+        else:
+            if recovered_offset > previous_offset:
+                self._retried = 0
+            self.offset = recovered_offset
+            self._do_request()
+
 
 class AsyncUploader(BaseUploader):
     def __init__(self, *args, **kwargs):
@@ -323,15 +329,21 @@ async def _do_request(self):
             await self._retry_or_cry(error)
 
     async def _retry_or_cry(self, error):
-        if self.retries > self._retried:
-            await asyncio.sleep(self.retry_delay)
-
-            self._retried += 1
-            try:
-                self.offset = self.get_offset()
-            except TusCommunicationError as err:
-                await self._retry_or_cry(err)
-            else:
-                await self._do_request()
-        else:
+        retry_attempt = self._retried
+        if self._retry_limit() <= retry_attempt:
+            raise error
+        if not self._should_retry(error, retry_attempt):
             raise error
+
+        await asyncio.sleep(self._retry_delay_seconds(retry_attempt))
+        self._retried += 1
+        previous_offset = self.offset
+        try:
+            recovered_offset = self.get_offset()
+        except TusCommunicationError as err:
+            await self._retry_or_cry(err)
+        else:
+            if recovered_offset > previous_offset:
+                self._retried = 0
+            self.offset = recovered_offset
+            await self._do_request()
