Fix calls to upload_chunk (#92)

Closes #82
tus · Dec 13, 2023 · 43b9b57 · 43b9b57
1 parent fd8db8f
commit 43b9b57
Show file tree

Hide file tree

Showing 3 changed files with 57 additions and 0 deletions.
diff --git a/tests/test_async_uploader.py b/tests/test_async_uploader.py
@@ -46,6 +46,28 @@ def test_upload_chunk(self):
             request_length = self.async_uploader.get_request_length()
             self.loop.run_until_complete(self.async_uploader.upload_chunk())
             self.assertEqual(self.async_uploader.offset, request_length)
+
+    def test_upload_chunk_with_creation(self):
+        with aioresponses() as resps:
+            resps.post(
+                self.client.url, status=201,
+                headers={
+                    "location": f"{self.client.url}hello"
+                }
+            )
+            resps.patch(
+                f"{self.client.url}hello",
+                headers={
+                    "upload-offset": "5"
+                }
+            )
+
+            uploader = self.client.async_uploader(
+                file_stream=io.BytesIO(b"hello")
+            )
+            self.loop.run_until_complete(uploader.upload_chunk())
+
+            self.assertEqual(uploader.url, f"{self.client.url}hello")
 
     def test_upload(self):
         with aioresponses() as resps:

diff --git a/tests/test_uploader.py b/tests/test_uploader.py
@@ -134,6 +134,29 @@ def test_upload_chunk(self, request_mock):
         self.uploader.upload_chunk()
         self.assertEqual(self.uploader.offset, request_length)
 
+    @responses.activate
+    def test_upload_chunk_with_creation(self):
+        responses.add(
+            responses.POST, self.client.url,
+            adding_headers={
+                "location": f"{self.client.url}hello"
+            }
+        )
+        responses.add(
+            responses.PATCH,
+            f"{self.client.url}hello",
+            adding_headers={
+                "upload-offset": "5"
+            }
+        )
+
+        uploader = self.client.uploader(
+            file_stream=io.BytesIO(b"hello")
+        )
+        uploader.upload_chunk()
+
+        self.assertEqual(uploader.url, f"{self.client.url}hello")
+
     @mock.patch('tusclient.uploader.uploader.TusRequest')
     def test_upload(self, request_mock):
         self.mock_request(request_mock)

diff --git a/tusclient/uploader/uploader.py b/tusclient/uploader/uploader.py
@@ -50,6 +50,12 @@ def upload_chunk(self):
         """
         self._retried = 0
 
+        # Ensure that we have a URL, as this is behavior we allowed previously.
+        # See https://github.com/tus/tus-py-client/issues/82.
+        if not self.url:
+            self.set_url(self.create_url())
+            self.offset = 0
+
         self._do_request()
         self.offset = int(self.request.response_headers.get("upload-offset"))
 
@@ -127,6 +133,12 @@ async def upload_chunk(self):
         """
         self._retried = 0
 
+        # Ensure that we have a URL, as this is behavior we allowed previously.
+        # See https://github.com/tus/tus-py-client/issues/82.
+        if not self.url:
+            self.set_url(await self.create_url())
+            self.offset = 0
+
         await self._do_request()
         self.offset = int(self.request.response_headers.get("upload-offset"))