From 4a196f7e30f4c9b2bf6f65f62044badcd25d71e5 Mon Sep 17 00:00:00 2001 From: Renaud Hartert Date: Wed, 13 Nov 2024 22:24:29 +0100 Subject: [PATCH] Add unit test for non-seekable stream --- tests/test_base_client.py | 36 ++++++++++++++++++++++++++++++++---- 1 file changed, 32 insertions(+), 4 deletions(-) diff --git a/tests/test_base_client.py b/tests/test_base_client.py index 4227737b..dcfc00af 100644 --- a/tests/test_base_client.py +++ b/tests/test_base_client.py @@ -317,10 +317,38 @@ def mock_iter_content(chunk_size): assert all(len(c) <= chunk_size for c in content_chunks) # chunks don't exceed size +def test_no_retry_on_non_seekable_stream(): + requests = [] + + # Always respond with a response that triggers a retry. + def inner(h: BaseHTTPRequestHandler): + content_length = int(h.headers.get('Content-Length', 0)) + if content_length > 0: + requests.append(h.rfile.read(content_length)) + + h.send_response(429) + h.send_header('Retry-After', '1') + h.end_headers() + + stream = io.BytesIO(b"test data") + stream.seekable = lambda: False # makes the stream appear non-seekable + + with http_fixture_server(inner) as host: + client = _BaseClient() + + # Should raise error immediately without retry. + with pytest.raises(DatabricksError): + client.do('POST', f'{host}/foo', data=stream) + + # Verify that only one request was made (no retries). + assert len(requests) == 1 + assert requests[0] == b"test data" + + def test_perform_resets_seekable_stream_on_error(): received_data = [] - # Response that triggers a retry. + # Always respond with a response that triggers a retry. def inner(h: BaseHTTPRequestHandler): content_length = int(h.headers.get('Content-Length', 0)) if content_length > 0: @@ -340,7 +368,7 @@ def inner(h: BaseHTTPRequestHandler): stream.read(4) assert stream.tell() == 4 - # Call perform which should fail but reset the stream. + # Should fail but reset the stream. with pytest.raises(DatabricksError): client._perform('POST', f'{host}/foo', data=stream) @@ -353,7 +381,7 @@ def inner(h: BaseHTTPRequestHandler): def test_perform_does_not_reset_nonseekable_stream_on_error(): received_data = [] - # Response that triggers a retry. + # Always respond with a response that triggers a retry. def inner(h: BaseHTTPRequestHandler): content_length = int(h.headers.get('Content-Length', 0)) if content_length > 0: @@ -374,7 +402,7 @@ def inner(h: BaseHTTPRequestHandler): stream.read(4) assert stream.tell() == 4 - # Call perform which should fail but reset the stream. + # Should fail without resetting the stream. with pytest.raises(DatabricksError): client._perform('POST', f'{host}/foo', data=stream)