Generate TUS Assembly helper

Author: kvz

tests/test_async_client.py

@@ -797,6 +797,32 @@ async def test_async_wait_for_assembly_rejects_non_json_poll_response(self):
                         f"{self.server.base_url}/assemblies/assembly-123"
                     )
 
+    async def test_async_create_tus_assembly_uses_contract_payload(self):
+        expected_response = Response(data={"ok": "ASSEMBLY_UPLOADING"}, status_code=200)
+
+        async with AsyncTransloadit("key", "secret", service=self.server.base_url) as client:
+            with mock.patch.object(
+                client,
+                "create_assembly",
+                new=mock.AsyncMock(return_value=expected_response),
+            ) as create_mock:
+                response = await client.create_tus_assembly(2)
+
+        self.assertIs(response, expected_response)
+        create_mock.assert_awaited_once_with(
+            data={
+                "await": False,
+                "steps": {
+                    ":original": {
+                        "output_meta": True,
+                        "result": "debug",
+                        "robot": "/upload/handle",
+                    },
+                },
+            },
+            extra_data={"num_expected_upload_files": 2},
+        )
+
     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")

tests/test_client.py

@@ -110,6 +110,29 @@ def test_wait_for_assembly_rejects_non_json_poll_response(self):
             with self.assertRaises(RuntimeError):
                 self.transloadit.wait_for_assembly("https://api2.example/assemblies/assembly-123")
 
+    def test_create_tus_assembly_uses_contract_payload(self):
+        expected_response = Response(data={"ok": "ASSEMBLY_UPLOADING"}, status_code=200)
+
+        with mock.patch.object(
+            self.transloadit, "create_assembly", return_value=expected_response
+        ) as create_mock:
+            response = self.transloadit.create_tus_assembly(2)
+
+        self.assertIs(response, expected_response)
+        create_mock.assert_called_once_with(
+            data={
+                "await": False,
+                "steps": {
+                    ":original": {
+                        "output_meta": True,
+                        "result": "debug",
+                        "robot": "/upload/handle",
+                    },
+                },
+            },
+            extra_data={"num_expected_upload_files": 2},
+        )
+
     def test_quotes_path_ids(self):
         with mock.patch.object(self.transloadit.request, 'get') as get_mock:
             self.transloadit.get_assembly(assembly_id='assembly/with?chars')

transloadit/async_client.py

@@ -245,6 +245,28 @@ async def update_template_credentials(self, identifier: str, data: Optional[dict
     # please report the issue instead of editing this block by hand; the source fix
     # belongs in the contract generator so all SDKs stay in sync.
 
+    async def create_tus_assembly(self, file_count: int):
+        """
+        Create a TUS-ready Assembly that waits for the requested number of resumable uploads before execution continues.
+        """
+        assembly = await self.create_assembly(
+            data={
+                "await": False,
+                "steps": {
+                    ":original": {
+                        "output_meta": True,
+                        "result": "debug",
+                        "robot": "/upload/handle",
+                    },
+                },
+            },
+            extra_data={
+                "num_expected_upload_files": file_count,
+            },
+        )
+
+        return assembly
+
     async def wait_for_assembly(self, assembly_url: str):
         """
         Wait for an Assembly to finish uploading and executing.

transloadit/client.py

@@ -253,6 +253,28 @@ def update_template_credentials(self, identifier: str, data: Optional[dict] = No
     # please report the issue instead of editing this block by hand; the source fix
     # belongs in the contract generator so all SDKs stay in sync.
 
+    def create_tus_assembly(self, file_count: int):
+        """
+        Create a TUS-ready Assembly that waits for the requested number of resumable uploads before execution continues.
+        """
+        assembly = self.create_assembly(
+            data={
+                "await": False,
+                "steps": {
+                    ":original": {
+                        "output_meta": True,
+                        "result": "debug",
+                        "robot": "/upload/handle",
+                    },
+                },
+            },
+            extra_data={
+                "num_expected_upload_files": file_count,
+            },
+        )
+
+        return assembly
+
     def wait_for_assembly(self, assembly_url: str):
         """
         Wait for an Assembly to finish uploading and executing.