Skip to content
Closed
Show file tree
Hide file tree
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 6 additions & 0 deletions databricks/sdk/service/files.py

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

127 changes: 125 additions & 2 deletions tests/test_files.py
Original file line number Diff line number Diff line change
Expand Up @@ -1279,8 +1279,6 @@ def run(self, config: Config):

session = requests.Session()
with requests_mock.Mocker(session=session) as session_mock:
session_mock.get(f"http://localhost/api/2.0/fs/files{MultipartUploadTestCase.path}", status_code=200)
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

cleanup


upload_state = SingleShotUploadState()

def custom_matcher(request):
Expand Down Expand Up @@ -1855,3 +1853,128 @@ def to_string(test_case):
)
def test_resumable_upload(config: Config, test_case: ResumableUploadTestCase):
test_case.run(config)


class SingleShotUploadContentLengthTestCase:

def __init__(
self,
name: str,
contents: Callable[[], io.IOBase],
expected_content_length: Optional[int],
cleanup: Callable[[io.IOBase], None] = None,
):
super().__init__()
self.name = name
self.contents = contents
self.expected_content_length = expected_content_length
self.cleanup = cleanup

def __str__(self):
return self.name

@staticmethod
def to_string(test_case):
return str(test_case)

def run(self, config: Config):
config = config.copy()
config.enable_experimental_files_api_client = False # enforce single-shot upload

file_path = "/test.txt"
contents = self.contents()

try:
with requests_mock.Mocker() as session_mock:

def custom_matcher(request):
request_url = urlparse(request.url)

if (
request_url.hostname == "localhost"
and request_url.path == f"/api/2.0/fs/files{file_path}"
and request.method == "PUT"
):
body = request.body.read()

if self.expected_content_length:
content_length = request.headers["Content-Length"]
assert self.expected_content_length == int(content_length)
assert len(body) == int(content_length)
else:
assert request.headers.get("Content-Length") is None

resp = requests.Response()
resp.status_code = 204
resp.request = request
resp._content = b""
return resp
return None

session_mock.add_matcher(matcher=custom_matcher)

w = WorkspaceClient(config=config)
w.files.upload(file_path, contents)
finally:
if self.cleanup:
self.cleanup(contents)


def make_non_seekable(stream: io.IOBase, disable_seek: bool = False, disable_tell: bool = False):
def raise_(ex):
raise ex

stream.seekable = lambda: False # checked by BaseClient._is_seekable_stream()

# requests.super_len() does not check seekable(), it calls seek() and tell() directly
if disable_seek:
stream.seek = lambda offset, whence: raise_(OSError())
if disable_tell:
stream.tell = lambda: raise_(OSError())
return stream


def create_file_stream(length: int) -> io.IOBase:
fd, temp_file = mkstemp()
with open(fd, "wb") as f:
f.write(os.urandom(length))

stream = open(temp_file, "rb")
stream.delete_temp_file = lambda: os.remove(temp_file)
return stream


@pytest.mark.parametrize(
"test_case",
[
SingleShotUploadContentLengthTestCase(
"Empty contents treated as unknown length", contents=lambda: io.BytesIO(b""), expected_content_length=None
),
SingleShotUploadContentLengthTestCase("Bytes", contents=lambda: io.BytesIO(b"abc"), expected_content_length=3),
SingleShotUploadContentLengthTestCase(
"seek disabled: length unknown",
contents=lambda: make_non_seekable(io.BytesIO(b"abc"), disable_seek=True),
expected_content_length=None,
),
SingleShotUploadContentLengthTestCase(
"tell disabled: length unknown",
contents=lambda: make_non_seekable(io.BytesIO(b"abc"), disable_tell=True),
expected_content_length=None,
),
SingleShotUploadContentLengthTestCase(
"File stream: length reported",
contents=lambda: create_file_stream(566),
expected_content_length=566,
cleanup=lambda stream: stream.delete_temp_file(),
),
SingleShotUploadContentLengthTestCase(
"File stream with tell disabled: length unknown",
contents=lambda: make_non_seekable(create_file_stream(239), disable_tell=True),
expected_content_length=None,
cleanup=lambda stream: stream.delete_temp_file(),
),
],
ids=SingleShotUploadContentLengthTestCase.to_string,
)
def test_content_length(config: Config, test_case: SingleShotUploadContentLengthTestCase):
test_case.run(config)
Loading