Improve async coverage and retry safety

Author: kvz

tests/test_async_client.py

@@ -292,8 +292,7 @@ async def test_async_client_methods_and_context_manager(self):
             self.assertEqual(response.data["ok"], "TEMPLATE_CREATED")
             self.assertEqual(response.data["template_name"], "foo")
 
-        self.assertIsNotNone(client.request.session)
-        self.assertTrue(client.request.session.closed)
+        self.assertIsNone(client.request.session)
 
         self.assertGreaterEqual(len(self.server.requests), 7)
         first_request = self.server.requests[0]
@@ -365,13 +364,28 @@ async def test_async_client_close_reopens_owned_session(self):
 
         await client.close()
         self.assertTrue(first_session.closed)
+        self.assertIsNone(client.request.session)
 
         second_session = await client.request._ensure_session()
         self.assertIsNot(first_session, second_session)
         self.assertFalse(second_session.closed)
 
         await client.close()
 
+    async def test_async_client_reopens_owned_session_when_session_is_closed(self):
+        client = AsyncTransloadit("key", "secret", service=self.server.base_url)
+
+        first_session = await client.request._ensure_session()
+        self.assertFalse(first_session.closed)
+
+        await first_session.close()
+        reopened_session = await client.request._ensure_session()
+
+        self.assertIsNot(first_session, reopened_session)
+        self.assertFalse(reopened_session.closed)
+
+        await client.close()
+
     async def test_async_client_delete_template_get_bill_and_plain_text_fallback(self):
         async with AsyncTransloadit("key", "secret", service=self.server.base_url) as client:
             response = await client.delete_template("tpl-1")
@@ -408,6 +422,22 @@ async def test_async_assembly_create_raises_on_plain_text_error_response(self):
 
         post_mock.assert_awaited_once()
 
+    async def test_async_assembly_create_returns_plain_text_success_response(self):
+        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()
+
+            with mock.patch.object(client.request, "post", new=mock.AsyncMock(return_value=plain_response)) as post_mock:
+                response = await assembly.create(wait=False, resumable=False)
+
+        self.assertIs(response, plain_response)
+        post_mock.assert_awaited_once()
+
     async def test_async_assembly_wait_raises_on_plain_text_poll_response(self):
         initial_response = Response(
             data={
@@ -439,6 +469,37 @@ async def test_async_assembly_wait_raises_on_plain_text_poll_response(self):
         )
         sleep_mock.assert_awaited_once_with(0)
 
+    async def test_async_assembly_wait_returns_plain_text_poll_response(self):
+        initial_response = 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"},
+        )
+        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()
+
+            with mock.patch.object(client.request, "post", new=mock.AsyncMock(return_value=initial_response)) as post_mock:
+                with mock.patch.object(client, "get_assembly", new=mock.AsyncMock(return_value=plain_response)) as get_mock:
+                    with mock.patch("asyncio.sleep", new_callable=mock.AsyncMock) as sleep_mock:
+                        response = await assembly.create(wait=True, resumable=False)
+
+        self.assertIs(response, plain_response)
+        post_mock.assert_awaited_once()
+        get_mock.assert_awaited_once_with(
+            assembly_url=f"{self.server.base_url}/assemblies/assembly-123"
+        )
+        sleep_mock.assert_awaited_once_with(0)
+
     def test_async_signed_smart_cdn_url_matches_sync_and_rejects_bad_types(self):
         async_client = AsyncTransloadit("test-key", "test-secret")
         sync_client = Transloadit("test-key", "test-secret")
@@ -633,7 +694,7 @@ def upload(self_inner):
         async with AsyncTransloadit("key", "secret", service=self.server.base_url) as client:
             assembly = client.new_assembly()
             upload = io.BytesIO(b"payload")
-            upload.name = "payload.bin"
+            upload.name = b"payload.bin"
             assembly.add_file(upload)
 
             rate_limited = Response(
@@ -730,6 +791,38 @@ def uploader(self, **kwargs):
         self.assertEqual(calls[0], ("client", f"{self.server.base_url}/uploads"))
         self.assertEqual(calls[1][0], "uploader")
 
+    async def test_async_assembly_non_resumable_rate_limit_raises_when_stream_cannot_be_snapshotted(self):
+        class _NonSeekableStream(io.BytesIO):
+            def tell(self):
+                raise OSError("tell failed")
+
+        reads = []
+
+        async def fake_post(path, data=None, extra_data=None, files=None):
+            file_stream = files["file"]
+            reads.append(file_stream.read())
+            return Response(
+                data={
+                    "error": "RATE_LIMIT_REACHED",
+                    "info": {"retryIn": 0},
+                },
+                status_code=200,
+                headers={},
+            )
+
+        async with AsyncTransloadit("key", "secret", service=self.server.base_url) as client:
+            assembly = client.new_assembly()
+            assembly.add_file(_NonSeekableStream(b"payload"))
+
+            with mock.patch.object(client.request, "post", new=mock.AsyncMock(side_effect=fake_post)) as post_mock:
+                with mock.patch("asyncio.sleep", new_callable=mock.AsyncMock) as sleep_mock:
+                    with self.assertRaises(RuntimeError):
+                        await assembly.create(resumable=False, retries=1)
+
+        self.assertEqual(reads, [b"payload"])
+        post_mock.assert_awaited_once()
+        sleep_mock.assert_not_awaited()
+
     async def test_async_assembly_resumable_rate_limit_returns_response_without_upload_when_retries_exhausted(self):
         calls = []
 
@@ -1169,9 +1262,10 @@ def test_async_assembly_helpers_cover_duplicate_names_and_rewind_edges(self):
 
         first.read(1)
         second.read(2)
-        positions = assembly._snapshot_file_positions()
+        positions, missing = assembly._snapshot_file_positions()
         self.assertEqual(positions["file"], 1)
         self.assertEqual(positions["file_1"], 2)
+        self.assertEqual(missing, [])
 
         first.read(1)
         second.read(1)
@@ -1181,8 +1275,9 @@ def test_async_assembly_helpers_cover_duplicate_names_and_rewind_edges(self):
 
         broken = _BrokenStream()
         assembly.files["broken"] = broken
-        positions = assembly._snapshot_file_positions()
+        positions, missing = assembly._snapshot_file_positions()
         self.assertNotIn("broken", positions)
+        self.assertEqual(missing, ["broken"])
 
         assembly._rewind_files({"missing": 4})
         with self.assertRaises(RuntimeError):

tests/test_response.py

@@ -1,7 +1,7 @@
 import unittest
 from unittest import mock
 
-from transloadit.response import Response
+from transloadit.response import Response, _MISSING
 
 
 class ResponseTest(unittest.TestCase):
@@ -21,6 +21,19 @@ def test_response_data_is_assignable_and_eager_for_sync_responses(self):
         self.assertEqual(response.status_code, 200)
         self.assertEqual(response.headers, {"X-Test": "1"})
 
+    def test_response_lazily_rehydrates_data_when_missing(self):
+        raw = mock.Mock()
+        raw.json.return_value = {"ok": "lazy"}
+        raw.status_code = 204
+        raw.headers = {"X-Test": "1"}
+
+        response = Response()
+        response._response = raw
+        response._data = _MISSING
+
+        self.assertEqual(response.data, {"ok": "lazy"})
+        raw.json.assert_called_once()
+
     def test_response_supports_async_preloaded_values_and_empty_default(self):
         empty = Response()
         self.assertIsNone(empty.data)

transloadit/async_assembly.py

@@ -43,12 +43,13 @@ def remove_file(self, field_name):
 
     def _snapshot_file_positions(self):
         positions = {}
+        missing = []
         for key, file_stream in self.files.items():
             try:
                 positions[key] = file_stream.tell()
             except (AttributeError, OSError, ValueError):
-                continue
-        return positions
+                missing.append(key)
+        return positions, missing
 
     def _rewind_files(self, positions):
         for key, position in positions.items():
@@ -84,7 +85,7 @@ async def create(self, wait=False, resumable=True, retries=3):
         Save/Submit the assembly for processing.
         """
         data = self.get_options()
-        file_positions = self._snapshot_file_positions()
+        file_positions, missing_file_positions = self._snapshot_file_positions()
         tus_retries = retries
         poll_retries = retries
 
@@ -107,6 +108,12 @@ async def create(self, wait=False, resumable=True, retries=3):
 
             if self._rate_limit_reached(response_data):
                 if retries:
+                    if not resumable and missing_file_positions:
+                        missing = ", ".join(repr(key) for key in missing_file_positions)
+                        raise RuntimeError(
+                            "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)
@@ -170,4 +177,7 @@ 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") == "RATE_LIMIT_REACHED"
+        return response_data.get("error") in {
+            "RATE_LIMIT_REACHED",
+            "ASSEMBLY_STATUS_FETCHING_RATE_LIMIT_REACHED",
+        }

transloadit/async_request.py

@@ -5,6 +5,7 @@
 import hashlib
 import hmac
 import json
+from types import MappingProxyType
 from datetime import datetime, timedelta, timezone
 
 import aiohttp
@@ -17,7 +18,10 @@
 
 def _get_upload_filename(file_stream, fallback):
     name = getattr(file_stream, "name", None)
-    if isinstance(name, (str, bytes, os.PathLike)):
+    if isinstance(name, (bytes, os.PathLike)):
+        name = os.fsdecode(name)
+
+    if isinstance(name, str):
         filename = os.path.basename(name)
         if filename:
             return filename
@@ -29,13 +33,13 @@ class AsyncRequest:
     Transloadit tailored asynchronous HTTP request object.
     """
 
-    HEADERS = {"Transloadit-Client": "python-sdk:" + __version__}
+    HEADERS = MappingProxyType({"Transloadit-Client": "python-sdk:" + __version__})
 
     def __init__(self, transloadit, session=None):
         self.transloadit = transloadit
         self._session = session
         self._owns_session = session is None
-        self._session_lock = asyncio.Lock()
+        self._session_lock = None
 
     @property
     def session(self):
@@ -44,8 +48,16 @@ def session(self):
     def _headers(self):
         return dict(self.HEADERS)
 
+    def _get_session_lock(self):
+        if self._session_lock is None:
+            # Create the lock lazily so the client can be instantiated before the loop starts.
+            self._session_lock = asyncio.Lock()
+        return self._session_lock
+
     async def _ensure_session(self):
-        async with self._session_lock:
+        if self._session is not None and not self._session.closed:
+            return self._session
+        async with self._get_session_lock():
             if self._session is None:
                 self._session = aiohttp.ClientSession()
                 self._owns_session = True
@@ -57,9 +69,10 @@ async def _ensure_session(self):
             return self._session
 
     async def aclose(self):
-        async with self._session_lock:
+        async with self._get_session_lock():
             if self._session is not None and not self._session.closed and self._owns_session:
                 await self._session.close()
+                self._session = None
 
     def _timeout(self, files=False):
         return aiohttp.ClientTimeout(