Add integration test for baremetal API

Author: nickles-lee

databricks/sdk/mixins/files.py

@@ -724,7 +724,7 @@ def _files_multipart_upload_complete(self, file_path: str, session_token: str, e
             body={'etags': etags})
         return
 
-    def _ps_multipart_upload_create_part_urls(self, session_token: str, page_token: str, page_size: int):
+    def multipart_upload_create_part_urls(self, session_token: str, *, page_token: str, page_size: int) -> MultipartUploadCreatePartUrlsResponse:
         """Request a set of presigned URLs for uploading parts of a file in a multipart upload session."""
 
         headers = {'Accept': 'application/json', 'Content-Type': 'application/json'}
@@ -762,7 +762,9 @@ def execute_presigned_url_request(self,
                 "Not all client-provided headers are populated") # TODO: Move to a dedicated exception type
 
         request_headers = {**presigned_url.headers, **headers}
-        self._api.do(presigned_url.method, presigned_url.url, headers=request_headers, data=data)
+        resp_headers = {}
+        resp = self._api.do(presigned_url.method, presigned_url.url, headers=request_headers, data=data, response_headers = resp_headers)
+        return (resp, resp_headers)
 
 
 class PresignedUrl:

tests/integration/test_files.py

@@ -8,6 +8,7 @@
 import pytest
 
 from databricks.sdk.core import DatabricksError
+from databricks.sdk.mixins.files import MultipartUploadCreate, MultipartUploadCreatePartUrlsResponse
 from databricks.sdk.service.catalog import VolumeType
 
 
@@ -228,6 +229,35 @@ def test_files_api_upload_download(ucws, random):
             with w.files.download(target_file).contents as f:
                 assert f.read() == b"some text data"
 
+def test_files_multipart_upload_download_baremetal(ucws, random):
+    w = ucws
+    schema = 'filesit-' + random()
+    volume = 'filesit-' + random()
+    with ResourceWithCleanup.create_schema(w, 'main', schema):
+        with ResourceWithCleanup.create_volume(w, 'main', schema, volume):
+            f = io.BytesIO(b"some text data")
+            target_file = f'/Volumes/main/{schema}/{volume}/filesit-with-?-and-#-{random()}.txt'
+
+            ## Create a session, upload the file in a single part, and then finalize the session
+            session: MultipartUploadCreate = w.files.multipart_upload_create(target_file)
+
+            url_page_one: MultipartUploadCreatePartUrlsResponse = w.files.multipart_upload_create_part_urls(session.session_token, page_size=1)
+
+            # Verify that a second page of URLs may be requested
+            w.files.multipart_upload_create_part_urls(session.session_token, page_token=url_page_one.next_page_token, page_size=1)
+
+
+            (part_upload_resp, part_upload_resp_headers) = w.files.execute_presigned_url_request(url_page_one.upload_part_urls[0], data=f)
+            print(part_upload_resp)
+            print(part_upload_resp_headers)
+
+            w.files.multipart_upload_complete(target_file, session.session_token, [part_upload_resp_headers["etag"]])
+
+            ## Download the file & assert that it's what we expect
+            with w.files.download(target_file).contents as f:
+                assert f.read() == b"some text data"
+
+
 
 def test_files_api_read_twice_from_one_download(ucws, random):
     w = ucws