Harden async upload edge cases

Author: kvz

tests/test_async_client.py

@@ -1,12 +1,14 @@
-import io
 import asyncio
+import io
 import json
+import threading
 from pathlib import Path
 from unittest import IsolatedAsyncioTestCase, mock
 
-from aiohttp import web
+from aiohttp import payload, web
 
 from transloadit.async_client import AsyncTransloadit
+from transloadit.async_request import _NonClosingUploadStream
 from transloadit.client import Transloadit
 from transloadit.response import Response
 
@@ -232,10 +234,22 @@ def __init__(self, payload):
         self.closed = False
         self.payload = payload
 
+    def delete(self, url, **kwargs):
+        self.calls.append((url, kwargs))
+        return _FakeResponseContext(self.payload)
+
+    def get(self, url, **kwargs):
+        self.calls.append((url, kwargs))
+        return _FakeResponseContext(self.payload)
+
     def post(self, url, **kwargs):
         self.calls.append((url, kwargs))
         return _FakeResponseContext(self.payload)
 
+    def put(self, url, **kwargs):
+        self.calls.append((url, kwargs))
+        return _FakeResponseContext(self.payload)
+
     async def close(self):
         self.closed = True
 
@@ -274,8 +288,8 @@ async def test_async_client_methods_and_context_manager(self):
             self.assertEqual(response.data["ok"], "ASSEMBLY_COMPLETED")
             self.assertEqual(response.data["assembly_id"], "abc123")
             self.assertEqual(response.status_code, 200)
-            self.assertIs(type(response.headers), dict)
             self.assertEqual(response.headers["X-Async-Route"], "get_assembly")
+            self.assertEqual(response.headers["x-async-route"], "get_assembly")
 
             response = await client.list_assemblies()
             self.assertEqual(response.data["items"], [])
@@ -368,6 +382,28 @@ async def test_async_client_normalizes_service_and_rejects_missing_ids(self):
         with self.assertRaises(RuntimeError):
             await closed_client.get_assembly(assembly_id="abc123")
 
+    async def test_async_client_quotes_path_ids(self):
+        session = _RecordingSession({"ok": "ASSEMBLY_COMPLETED"})
+        client = AsyncTransloadit("key", "secret", service=self.server.base_url, session=session)
+
+        await client.get_assembly(assembly_id="assembly/with?chars")
+        await client.cancel_assembly(assembly_id="cancel/with?chars")
+        await client.get_template("template/with?chars")
+        await client.update_template("update/with?chars", {"name": "foo"})
+        await client.delete_template("delete/with?chars")
+
+        urls = [call[0] for call in session.calls]
+        self.assertEqual(
+            urls,
+            [
+                f"{self.server.base_url}/assemblies/assembly%2Fwith%3Fchars",
+                f"{self.server.base_url}/assemblies/cancel%2Fwith%3Fchars",
+                f"{self.server.base_url}/templates/template%2Fwith%3Fchars",
+                f"{self.server.base_url}/templates/update%2Fwith%3Fchars",
+                f"{self.server.base_url}/templates/delete%2Fwith%3Fchars",
+            ],
+        )
+
     async def test_async_client_close_reopens_owned_session(self):
         client = AsyncTransloadit("key", "secret", service=self.server.base_url)
 
@@ -469,6 +505,34 @@ async def test_async_assembly_create_returns_plain_text_success_response(self):
         self.assertIs(response, plain_response)
         post_mock.assert_awaited_once()
 
+    async def test_async_assembly_resumable_plain_text_success_response_raises_before_tus_upload(self):
+        calls = []
+
+        class _TusClient:
+            def __init__(self, tus_url):
+                calls.append(("client", tus_url))
+
+            def uploader(self, **kwargs):
+                raise AssertionError("TUS upload should not start without upload URLs")
+
+        plain_response = Response(
+            data="plain assembly response",
+            status_code=200,
+            headers={"X-Async-Route": "plain"},
+        )
+
+        async with AsyncTransloadit("key", "secret", service=self.server.base_url) as client:
+            assembly = client.new_assembly()
+            assembly.add_file(io.BytesIO(b"payload"))
+
+            with mock.patch.object(client.request, "post", new=mock.AsyncMock(return_value=plain_response)) as post_mock:
+                with mock.patch("transloadit.async_assembly.tus.TusClient", new=_TusClient):
+                    with self.assertRaises(RuntimeError):
+                        await assembly.create(resumable=True)
+
+        post_mock.assert_awaited_once()
+        self.assertEqual(calls, [])
+
     async def test_async_assembly_wait_raises_on_plain_text_poll_response(self):
         initial_response = Response(
             data={
@@ -1182,6 +1246,7 @@ async def test_async_assembly_wait_returns_last_poll_response_when_budget_exhaus
 
     async def test_async_assembly_non_resumable_rate_limit_rewinds_files_for_retry(self):
         reads = []
+        upload = io.BytesIO(b"payload")
 
         async def fake_post(path, data=None, extra_data=None, files=None):
             file_stream = files["file"]
@@ -1201,12 +1266,15 @@ async def fake_post(path, data=None, extra_data=None, files=None):
                 headers={},
             )
 
+        async def fake_sleep(delay):
+            self.assertEqual(upload.tell(), 0)
+
         async with AsyncTransloadit("key", "secret", service=self.server.base_url) as client:
             assembly = client.new_assembly()
-            assembly.add_file(io.BytesIO(b"payload"))
+            assembly.add_file(upload)
 
             with mock.patch.object(client.request, "post", new=mock.AsyncMock(side_effect=fake_post)):
-                with mock.patch("asyncio.sleep", new_callable=mock.AsyncMock) as sleep_mock:
+                with mock.patch("asyncio.sleep", new=mock.AsyncMock(side_effect=fake_sleep)) as sleep_mock:
                     response = await assembly.create(resumable=False, retries=2)
 
         self.assertEqual(response.data["ok"], "ASSEMBLY_COMPLETED")
@@ -1238,6 +1306,47 @@ def seek(self, position, *args, **kwargs):
 
         post_mock.assert_awaited_once()
 
+    async def test_async_assembly_rate_limit_ignores_malformed_error_values(self):
+        client = AsyncTransloadit("key", "secret", service=self.server.base_url)
+        assembly = client.new_assembly()
+
+        self.assertFalse(assembly._rate_limit_reached({"error": ["RATE_LIMIT_REACHED"]}))
+        self.assertFalse(assembly._rate_limit_reached({"error": {"code": "RATE_LIMIT_REACHED"}}))
+
+    async def test_async_tus_upload_cancellation_waits_for_thread_to_finish(self):
+        client = AsyncTransloadit("key", "secret", service=self.server.base_url)
+        assembly = client.new_assembly()
+        started = threading.Event()
+        release = threading.Event()
+        finished = threading.Event()
+
+        def blocking_upload(assembly_url, tus_url, retries):
+            started.set()
+            release.wait(timeout=5)
+            finished.set()
+
+        assembly._do_tus_upload = blocking_upload
+        upload_task = asyncio.create_task(
+            assembly._do_tus_upload_async(
+                f"{self.server.base_url}/assemblies/assembly-123",
+                f"{self.server.base_url}/uploads",
+                retries=1,
+            )
+        )
+
+        await asyncio.to_thread(started.wait, 5)
+        upload_task.cancel()
+        await asyncio.sleep(0.05)
+
+        self.assertFalse(upload_task.done())
+        self.assertFalse(finished.is_set())
+
+        release.set()
+        with self.assertRaises(asyncio.CancelledError):
+            await upload_task
+
+        self.assertTrue(finished.is_set())
+
     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)
@@ -1250,9 +1359,25 @@ async def test_async_request_uses_connect_and_read_timeouts_for_uploads(self):
         timeout = session.calls[0][1]["timeout"]
         self.assertIsNone(timeout.total)
         self.assertEqual(timeout.sock_connect, 60)
-        self.assertIsNone(timeout.sock_read)
+        self.assertEqual(timeout.sock_read, 60)
         self.assertEqual(session.calls[0][1]["data"]._fields[2][1]["Content-Type"], "image/jpeg")
 
+    async def test_async_request_upload_does_not_close_caller_stream(self):
+        fixture_path = Path(__file__).resolve().parents[1] / "LICENSE"
+        upload = fixture_path.open("rb")
+
+        try:
+            upload_payload = payload.get_payload(_NonClosingUploadStream(upload))
+            await upload_payload.close()
+            await asyncio.sleep(0.05)
+
+            self.assertFalse(upload.closed)
+            upload.seek(0)
+            self.assertEqual(upload.read(5), fixture_path.read_bytes()[:5])
+        finally:
+            if not upload.closed:
+                upload.close()
+
     async def test_async_request_payload_preserves_custom_auth_constraints(self):
         client = AsyncTransloadit("key", "secret", service=self.server.base_url)
 

transloadit/async_assembly.py

@@ -78,7 +78,16 @@ def _do_tus_upload(self, assembly_url, tus_url, retries):
             ).upload()
 
     async def _do_tus_upload_async(self, assembly_url, tus_url, retries):
-        await asyncio.to_thread(self._do_tus_upload, assembly_url, tus_url, retries)
+        upload_task = asyncio.create_task(
+            asyncio.to_thread(self._do_tus_upload, assembly_url, tus_url, retries)
+        )
+        try:
+            await asyncio.shield(upload_task)
+        except asyncio.CancelledError:
+            try:
+                await asyncio.shield(upload_task)
+            finally:
+                raise
 
     async def create(self, wait=False, resumable=True, retries=3):
         """
@@ -104,6 +113,8 @@ async def create(self, wait=False, resumable=True, retries=3):
             if response_data is None:
                 if response.status_code >= 400:
                     raise RuntimeError(f"Unexpected non-JSON response ({response.status_code}).")
+                if resumable and self.files:
+                    raise RuntimeError("Resumable assembly response is missing upload URLs.")
                 return response
 
             if self._rate_limit_reached(response_data):
@@ -114,9 +125,9 @@ async def create(self, wait=False, resumable=True, retries=3):
                             "Cannot retry non-resumable upload because these file streams are not seekable: "
                             f"{missing}"
                         )
-                    await asyncio.sleep(response_data.get("info", {}).get("retryIn", 1))
                     if not resumable:
                         self._rewind_files(file_positions)
+                    await asyncio.sleep(response_data.get("info", {}).get("retryIn", 1))
                     retries -= 1
                     continue
                 return response
@@ -176,7 +187,8 @@ def _assembly_finished(self, response_data):
         return is_aborted or is_canceled or is_completed or (is_failed and not (is_fetch_rate_limit or is_submit_rate_limit))
 
     def _rate_limit_reached(self, response_data):
-        return response_data.get("error") in {
+        error = response_data.get("error")
+        return isinstance(error, str) and error in {
             "RATE_LIMIT_REACHED",
             "ASSEMBLY_STATUS_FETCHING_RATE_LIMIT_REACHED",
         }

transloadit/async_client.py

@@ -2,7 +2,7 @@
 import hmac
 import time
 from typing import List, Optional, Union
-from urllib.parse import quote_plus, urlencode
+from urllib.parse import quote, quote_plus, urlencode
 
 from . import async_assembly, async_request, async_template
 
@@ -13,6 +13,10 @@ def _stringify_url_param(value: Union[str, int, float, bool]) -> str:
     return str(value)
 
 
+def _quote_path_segment(value: str) -> str:
+    return quote(str(value), safe="")
+
+
 class AsyncTransloadit:
     """
     Asynchronous client interface to the Transloadit API.
@@ -61,7 +65,7 @@ async def get_assembly(self, assembly_id: str = None, assembly_url: str = None):
         if not (assembly_id or assembly_url):
             raise ValueError("Either 'assembly_id' or 'assembly_url' cannot be None.")
 
-        url = assembly_url if assembly_url else f"/assemblies/{assembly_id}"
+        url = assembly_url if assembly_url else f"/assemblies/{_quote_path_segment(assembly_id)}"
         return await self.request.get(url)
 
     async def list_assemblies(self, params: dict = None):
@@ -77,14 +81,14 @@ async def cancel_assembly(self, assembly_id: str = None, assembly_url: str = Non
         if not (assembly_id or assembly_url):
             raise ValueError("Either 'assembly_id' or 'assembly_url' cannot be None.")
 
-        url = assembly_url if assembly_url else f"/assemblies/{assembly_id}"
+        url = assembly_url if assembly_url else f"/assemblies/{_quote_path_segment(assembly_id)}"
         return await self.request.delete(url)
 
     async def get_template(self, template_id: str):
         """
         Get the template specified by the 'template_id'.
         """
-        return await self.request.get(f"/templates/{template_id}")
+        return await self.request.get(f"/templates/{_quote_path_segment(template_id)}")
 
     async def list_templates(self, params: Optional[dict] = None):
         """
@@ -102,13 +106,13 @@ async def update_template(self, template_id: str, data: dict):
         """
         Update the template specified by the 'template_id'.
         """
-        return await self.request.put(f"/templates/{template_id}", data=data)
+        return await self.request.put(f"/templates/{_quote_path_segment(template_id)}", data=data)
 
     async def delete_template(self, template_id: str):
         """
         Delete the template specified by the 'template_id'.
         """
-        return await self.request.delete(f"/templates/{template_id}")
+        return await self.request.delete(f"/templates/{_quote_path_segment(template_id)}")
 
     async def get_bill(self, month: int, year: int):
         """