Skip to content

Commit 70e8ba9

Browse files
committed
Generate TUS Assembly helper
1 parent 24738de commit 70e8ba9

4 files changed

Lines changed: 93 additions & 0 deletions

File tree

tests/test_async_client.py

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -797,6 +797,32 @@ async def test_async_wait_for_assembly_rejects_non_json_poll_response(self):
797797
f"{self.server.base_url}/assemblies/assembly-123"
798798
)
799799

800+
async def test_async_create_tus_assembly_uses_contract_payload(self):
801+
expected_response = Response(data={"ok": "ASSEMBLY_UPLOADING"}, status_code=200)
802+
803+
async with AsyncTransloadit("key", "secret", service=self.server.base_url) as client:
804+
with mock.patch.object(
805+
client,
806+
"create_assembly",
807+
new=mock.AsyncMock(return_value=expected_response),
808+
) as create_mock:
809+
response = await client.create_tus_assembly(2)
810+
811+
self.assertIs(response, expected_response)
812+
create_mock.assert_awaited_once_with(
813+
data={
814+
"await": False,
815+
"steps": {
816+
":original": {
817+
"output_meta": True,
818+
"result": "debug",
819+
"robot": "/upload/handle",
820+
},
821+
},
822+
},
823+
extra_data={"num_expected_upload_files": 2},
824+
)
825+
800826
def test_async_signed_smart_cdn_url_matches_sync_and_rejects_bad_types(self):
801827
async_client = AsyncTransloadit("test-key", "test-secret")
802828
sync_client = Transloadit("test-key", "test-secret")

tests/test_client.py

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -110,6 +110,29 @@ def test_wait_for_assembly_rejects_non_json_poll_response(self):
110110
with self.assertRaises(RuntimeError):
111111
self.transloadit.wait_for_assembly("https://api2.example/assemblies/assembly-123")
112112

113+
def test_create_tus_assembly_uses_contract_payload(self):
114+
expected_response = Response(data={"ok": "ASSEMBLY_UPLOADING"}, status_code=200)
115+
116+
with mock.patch.object(
117+
self.transloadit, "create_assembly", return_value=expected_response
118+
) as create_mock:
119+
response = self.transloadit.create_tus_assembly(2)
120+
121+
self.assertIs(response, expected_response)
122+
create_mock.assert_called_once_with(
123+
data={
124+
"await": False,
125+
"steps": {
126+
":original": {
127+
"output_meta": True,
128+
"result": "debug",
129+
"robot": "/upload/handle",
130+
},
131+
},
132+
},
133+
extra_data={"num_expected_upload_files": 2},
134+
)
135+
113136
def test_quotes_path_ids(self):
114137
with mock.patch.object(self.transloadit.request, 'get') as get_mock:
115138
self.transloadit.get_assembly(assembly_id='assembly/with?chars')

transloadit/async_client.py

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -245,6 +245,28 @@ async def update_template_credentials(self, identifier: str, data: Optional[dict
245245
# please report the issue instead of editing this block by hand; the source fix
246246
# belongs in the contract generator so all SDKs stay in sync.
247247

248+
async def create_tus_assembly(self, file_count: int):
249+
"""
250+
Create a TUS-ready Assembly that waits for the requested number of resumable uploads before execution continues.
251+
"""
252+
assembly = await self.create_assembly(
253+
data={
254+
"await": False,
255+
"steps": {
256+
":original": {
257+
"output_meta": True,
258+
"result": "debug",
259+
"robot": "/upload/handle",
260+
},
261+
},
262+
},
263+
extra_data={
264+
"num_expected_upload_files": file_count,
265+
},
266+
)
267+
268+
return assembly
269+
248270
async def wait_for_assembly(self, assembly_url: str):
249271
"""
250272
Wait for an Assembly to finish uploading and executing.

transloadit/client.py

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -253,6 +253,28 @@ def update_template_credentials(self, identifier: str, data: Optional[dict] = No
253253
# please report the issue instead of editing this block by hand; the source fix
254254
# belongs in the contract generator so all SDKs stay in sync.
255255

256+
def create_tus_assembly(self, file_count: int):
257+
"""
258+
Create a TUS-ready Assembly that waits for the requested number of resumable uploads before execution continues.
259+
"""
260+
assembly = self.create_assembly(
261+
data={
262+
"await": False,
263+
"steps": {
264+
":original": {
265+
"output_meta": True,
266+
"result": "debug",
267+
"robot": "/upload/handle",
268+
},
269+
},
270+
},
271+
extra_data={
272+
"num_expected_upload_files": file_count,
273+
},
274+
)
275+
276+
return assembly
277+
256278
def wait_for_assembly(self, assembly_url: str):
257279
"""
258280
Wait for an Assembly to finish uploading and executing.

0 commit comments

Comments
 (0)