Refine async retries and upload metadata

Author: kvz

tests/test_async_client.py

@@ -628,7 +628,7 @@ def __init__(self, tus_url):
                 calls.append(("client", tus_url))
 
             def uploader(self, **kwargs):
-                calls.append(("uploader", kwargs["metadata"]))
+                calls.append(("uploader", kwargs["metadata"], kwargs["retries"]))
 
                 class _Uploader:
                     def upload(self_inner):
@@ -673,7 +673,7 @@ def upload(self_inner):
         self.assertEqual(post_mock.await_count, 2)
         self.assertEqual(to_thread_mock.await_count, 1)
         self.assertEqual(calls[0], ("client", f"{self.server.base_url}/uploads"))
-        self.assertEqual(calls[1][0], "uploader")
+        self.assertEqual(calls[1], ("uploader", {"assembly_url": f"{self.server.base_url}/assemblies/assembly-123", "fieldname": "file", "filename": "payload.bin"}, 2))
 
     async def test_async_assembly_resumable_rate_limit_skips_rewind_before_retrying(self):
         calls = []
@@ -868,7 +868,7 @@ def uploader(self, **kwargs):
         async with AsyncTransloadit("key", "secret", service=self.server.base_url) as client:
             assembly = client.new_assembly()
             upload = io.BytesIO(b"payload")
-            upload.name = None
+            upload.name = 123
             assembly.add_file(upload, "explicit_field")
 
             with mock.patch(
@@ -953,6 +953,46 @@ async def test_async_assembly_wait_retries_after_polling_rate_limit(self):
         )
         self.assertEqual(sleep_mock.await_args_list, [mock.call(0), mock.call(0)])
 
+    async def test_async_assembly_wait_returns_last_poll_response_when_budget_exhausted(self):
+        async with AsyncTransloadit("key", "secret", service=self.server.base_url) as client:
+            assembly = client.new_assembly()
+
+            initial = Response(
+                data={
+                    "ok": "ASSEMBLY_PROCESSING",
+                    "info": {"retryIn": 0},
+                    "assembly_ssl_url": f"{self.server.base_url}/assemblies/assembly-123",
+                },
+                status_code=200,
+                headers={"X-Async-Route": "initial"},
+            )
+            rate_limited = Response(
+                data={
+                    "ok": "ASSEMBLY_PROCESSING",
+                    "error": "RATE_LIMIT_REACHED",
+                    "info": {"retryIn": 0},
+                    "assembly_ssl_url": f"{self.server.base_url}/assemblies/assembly-123",
+                },
+                status_code=200,
+                headers={"X-Async-Route": "rate_limited"},
+            )
+
+            with mock.patch.object(client.request, "post", new=mock.AsyncMock(return_value=initial)) as post_mock:
+                with mock.patch.object(
+                    client,
+                    "get_assembly",
+                    new=mock.AsyncMock(return_value=rate_limited),
+                ) as get_mock:
+                    with mock.patch("asyncio.sleep", new_callable=mock.AsyncMock) as sleep_mock:
+                        response = await assembly.create(wait=True, resumable=False, retries=1)
+
+        self.assertEqual(response.data["error"], "RATE_LIMIT_REACHED")
+        post_mock.assert_awaited_once()
+        get_mock.assert_awaited_once_with(
+            assembly_url=f"{self.server.base_url}/assemblies/assembly-123"
+        )
+        self.assertEqual(sleep_mock.await_args_list, [mock.call(0)])
+
     async def test_async_assembly_non_resumable_rate_limit_rewinds_files_for_retry(self):
         reads = []
 
@@ -1015,7 +1055,7 @@ async def test_async_request_uses_connect_and_read_timeouts_for_uploads(self):
         session = _RecordingSession({"ok": "ASSEMBLY_COMPLETED"})
         client = AsyncTransloadit("key", "secret", service=self.server.base_url, session=session)
         upload = io.BytesIO(b"payload")
-        upload.name = None
+        upload.name = "clip.jpg"
 
         response = await client.request.post("/assemblies", data={"foo": "bar"}, files={"file": upload})
 
@@ -1024,6 +1064,18 @@ async def test_async_request_uses_connect_and_read_timeouts_for_uploads(self):
         self.assertIsNone(timeout.total)
         self.assertEqual(timeout.sock_connect, 60)
         self.assertIsNone(timeout.sock_read)
+        self.assertEqual(session.calls[0][1]["data"]._fields[2][1]["Content-Type"], "image/jpeg")
+
+    async def test_async_request_uses_filename_fallback_for_trailing_slash_stream_name(self):
+        session = _RecordingSession({"ok": "ASSEMBLY_COMPLETED"})
+        client = AsyncTransloadit("key", "secret", service=self.server.base_url, session=session)
+        upload = io.BytesIO(b"payload")
+        upload.name = "/tmp/"
+
+        response = await client.request.post("/assemblies", data={"foo": "bar"}, files={"file": upload})
+
+        self.assertEqual(response.data["ok"], "ASSEMBLY_COMPLETED")
+        self.assertEqual(session.calls[0][1]["data"]._fields[2][0]["filename"], "file")
 
     async def test_async_resumable_upload_uses_to_thread(self):
         calls = []

transloadit/async_assembly.py

@@ -1,9 +1,9 @@
 import asyncio
-import os
 
 from tusclient import client as tus
 
 from . import optionbuilder
+from .async_request import _get_upload_filename
 
 
 class AsyncAssembly(optionbuilder.OptionBuilder):
@@ -63,11 +63,11 @@ def _rewind_files(self, positions):
     def _do_tus_upload(self, assembly_url, tus_url, retries):
         tus_client = tus.TusClient(tus_url)
         for key, file_stream in self.files.items():
-            filename = getattr(file_stream, "name", None) or key
+            filename = _get_upload_filename(file_stream, key)
             metadata = {
                 "assembly_url": assembly_url,
                 "fieldname": key,
-                "filename": os.path.basename(filename) or key,
+                "filename": filename,
             }
             tus_client.uploader(
                 file_stream=file_stream,
@@ -85,55 +85,69 @@ async def create(self, wait=False, resumable=True, retries=3):
         """
         data = self.get_options()
         file_positions = self._snapshot_file_positions()
-        if resumable:
-            extra_data = {"tus_num_expected_upload_files": len(self.files)}
-            response = await self.transloadit.request.post(
-                "/assemblies", extra_data=extra_data, data=data
-            )
-        else:
-            response = await self.transloadit.request.post(
-                "/assemblies", data=data, files=self.files
-            )
-
-        response_data = self._response_data(response)
-        if response_data is None:
-            return response
+        tus_retries = retries
+        poll_retries = retries
+
+        while True:
+            if resumable:
+                extra_data = {"tus_num_expected_upload_files": len(self.files)} if self.files else None
+                response = await self.transloadit.request.post(
+                    "/assemblies", extra_data=extra_data, data=data
+                )
+            else:
+                response = await self.transloadit.request.post(
+                    "/assemblies", data=data, files=self.files
+                )
 
-        if self._rate_limit_reached(response_data):
-            if retries:
-                await asyncio.sleep(response_data.get("info", {}).get("retryIn", 1))
-                if not resumable:
-                    self._rewind_files(file_positions)
-                return await self.create(wait, resumable, retries - 1)
-            return response
+            response_data = self._response_data(response)
+            if response_data is None:
+                return response
 
-        error = response_data.get("error")
-        assembly_url = response_data.get("assembly_ssl_url")
-        tus_url = response_data.get("tus_url")
+            if self._rate_limit_reached(response_data):
+                if retries:
+                    await asyncio.sleep(response_data.get("info", {}).get("retryIn", 1))
+                    if not resumable:
+                        self._rewind_files(file_positions)
+                    retries -= 1
+                    continue
+                return response
 
-        if error is not None:
-            return response
+            error = response_data.get("error")
+            assembly_url = response_data.get("assembly_ssl_url")
+            tus_url = response_data.get("tus_url")
 
-        if resumable and self.files:
-            if not assembly_url or not tus_url:
+            if error is not None:
                 return response
-            await self._do_tus_upload_async(assembly_url, tus_url, retries)
 
-        if wait:
-            if not assembly_url:
-                return response
-            while not self._assembly_finished(response_data):
-                sleep_time = response_data.get("info", {}).get("retryIn", 1)
-                await asyncio.sleep(sleep_time)
-                response = await self.transloadit.get_assembly(
-                    assembly_url=assembly_url or response_data.get("assembly_ssl_url")
-                )
-                response_data = self._response_data(response)
-                if response_data is None:
+            if resumable and self.files:
+                if not assembly_url or not tus_url:
                     return response
-                assembly_url = response_data.get("assembly_ssl_url") or assembly_url
+                await self._do_tus_upload_async(assembly_url, tus_url, tus_retries)
 
-        return response
+            if wait:
+                if not assembly_url:
+                    return response
+
+                poll_response = response
+                poll_data = response_data
+                remaining_polls = poll_retries
+                while not self._assembly_finished(poll_data):
+                    if remaining_polls <= 0:
+                        return poll_response
+                    sleep_time = poll_data.get("info", {}).get("retryIn", 1)
+                    await asyncio.sleep(sleep_time)
+                    poll_response = await self.transloadit.get_assembly(
+                        assembly_url=assembly_url or poll_data.get("assembly_ssl_url")
+                    )
+                    poll_data = self._response_data(poll_response)
+                    if poll_data is None:
+                        return poll_response
+                    assembly_url = poll_data.get("assembly_ssl_url") or assembly_url
+                    remaining_polls -= 1
+
+                return poll_response
+
+            return response
 
     def _response_data(self, response):
         data = response.data

transloadit/async_request.py

@@ -1,4 +1,5 @@
 import asyncio
+import mimetypes
 import os
 import copy
 import hashlib
@@ -14,6 +15,15 @@
 TIMEOUT = 60
 
 
+def _get_upload_filename(file_stream, fallback):
+    name = getattr(file_stream, "name", None)
+    if isinstance(name, (str, bytes, os.PathLike)):
+        filename = os.path.basename(name)
+        if filename:
+            return filename
+    return fallback
+
+
 class AsyncRequest:
     """
     Transloadit tailored asynchronous HTTP request object.
@@ -96,8 +106,9 @@ async def post(self, path, data=None, extra_data=None, files=None):
                 form.add_field(key, value)
 
             for key, file_stream in files.items():
-                filename = os.path.basename(getattr(file_stream, "name", None) or key) or key
-                form.add_field(key, file_stream, filename=filename)
+                filename = _get_upload_filename(file_stream, key)
+                content_type = mimetypes.guess_type(filename)[0] or "application/octet-stream"
+                form.add_field(key, file_stream, filename=filename, content_type=content_type)
             payload = form
         else:
             payload = self._normalize_payload(data)