Handle plain-text async assembly responses

Author: kvz

tests/test_async_client.py

@@ -357,6 +357,21 @@ 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_close_reopens_owned_session(self):
+        client = AsyncTransloadit("key", "secret", service=self.server.base_url)
+
+        first_session = await client.request._ensure_session()
+        self.assertFalse(first_session.closed)
+
+        await client.close()
+        self.assertTrue(first_session.closed)
+
+        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_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")
@@ -377,6 +392,59 @@ async def test_async_client_delete_template_get_bill_and_plain_text_fallback(sel
         self.assertEqual(response.status_code, 200)
         self.assertEqual(response.headers["X-Async-Route"], "get_assembly_plain")
 
+    async def test_async_assembly_create_returns_plain_text_response_without_crashing(self):
+        plain_response = Response(
+            data="plain assembly response",
+            status_code=502,
+            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:
+                with mock.patch.object(client, "get_assembly", new=mock.AsyncMock()) 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)
+        self.assertEqual(response.data, "plain assembly response")
+        post_mock.assert_awaited_once()
+        get_mock.assert_not_awaited()
+        sleep_mock.assert_not_awaited()
+
+    async def test_async_assembly_wait_returns_plain_text_poll_response_without_crashing(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=502,
+            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)
+        self.assertEqual(response.data, "plain assembly 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")

transloadit/async_assembly.py

@@ -95,42 +95,53 @@ async def create(self, wait=False, resumable=True, retries=3):
                 "/assemblies", data=data, files=self.files
             )
 
-        if self._rate_limit_reached(response):
+        response_data = self._response_data(response)
+        if response_data is None:
+            return response
+
+        if self._rate_limit_reached(response_data):
             if retries:
-                await asyncio.sleep(response.data.get("info", {}).get("retryIn", 1))
+                await asyncio.sleep(response_data.get("info", {}).get("retryIn", 1))
                 self._rewind_files(file_positions)
                 return await self.create(wait, resumable, retries - 1)
             return response
 
         if resumable:
             await self._do_tus_upload_async(
-                response.data.get("assembly_ssl_url"),
-                response.data.get("tus_url"),
+                response_data.get("assembly_ssl_url"),
+                response_data.get("tus_url"),
                 retries,
             )
 
         if wait:
-            assembly_url = response.data.get("assembly_ssl_url")
-            while not self._assembly_finished(response):
-                sleep_time = response.data.get("info", {}).get("retryIn", 1)
+            assembly_url = response_data.get("assembly_ssl_url")
+            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")
+                    assembly_url=assembly_url or response_data.get("assembly_ssl_url")
                 )
-                assembly_url = response.data.get("assembly_ssl_url") or assembly_url
+                response_data = self._response_data(response)
+                if response_data is None:
+                    return response
+                assembly_url = response_data.get("assembly_ssl_url") or assembly_url
 
         return response
 
-    def _assembly_finished(self, response):
-        status = response.data.get("ok")
+    def _response_data(self, response):
+        data = response.data
+        return data if isinstance(data, dict) else None
+
+    def _assembly_finished(self, response_data):
+        status = response_data.get("ok")
         is_aborted = status == "REQUEST_ABORTED"
         is_canceled = status == "ASSEMBLY_CANCELED"
         is_completed = status == "ASSEMBLY_COMPLETED"
-        error = response.data.get("error")
+        error = response_data.get("error")
         is_failed = error is not None
         is_fetch_rate_limit = error == "ASSEMBLY_STATUS_FETCHING_RATE_LIMIT_REACHED"
         is_submit_rate_limit = error == "RATE_LIMIT_REACHED"
         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):
-        return response.data.get("error") == "RATE_LIMIT_REACHED"
+    def _rate_limit_reached(self, response_data):
+        return response_data.get("error") == "RATE_LIMIT_REACHED"

transloadit/async_request.py

@@ -44,8 +44,9 @@ async def _ensure_session(self):
             return self._session
 
     async def aclose(self):
-        if self._session is not None and not self._session.closed and self._owns_session:
-            await self._session.close()
+        async with self._session_lock:
+            if self._session is not None and not self._session.closed and self._owns_session:
+                await self._session.close()
 
     def _timeout(self, files=False):
         return aiohttp.ClientTimeout(