Make fmt

Author: renaudhartert-db

databricks/sdk/_base_client.py

@@ -51,7 +51,7 @@ def __init__(self,
                  extra_error_customizers: List[_ErrorCustomizer] = None,
                  debug_headers: bool = False,
                  clock: Clock = None,
-                 streaming_buffer_size: int = 1024 * 1024):  # 1MB
+                 streaming_buffer_size: int = 1024 * 1024): # 1MB
         """
         :param debug_truncate_bytes:
         :param retry_timeout_seconds:

tests/test_base_client.py

@@ -1,9 +1,9 @@
 from http.server import BaseHTTPRequestHandler
 from typing import Iterator, List
+from unittest.mock import Mock
 
 import pytest
 import requests
-from unittest.mock import Mock
 
 from databricks.sdk import errors, useragent
 from databricks.sdk._base_client import _BaseClient, _StreamingResponse
@@ -279,34 +279,36 @@ def inner(h: BaseHTTPRequestHandler):
     assert len(requests) == 2
 
 
-@pytest.mark.parametrize('chunk_size,expected_chunks,data_size', [
-    (5, 20, 100),    # 100 / 5 bytes per chunk = 20 chunks
-    (10, 10, 100),   # 100 / 10 bytes per chunk = 10 chunks
-    (200, 1, 100),   # 100 / 200 bytes per chunk = 1 chunk
-])
+@pytest.mark.parametrize('chunk_size,expected_chunks,data_size',
+                         [(5, 20, 100), # 100 / 5 bytes per chunk = 20 chunks
+                          (10, 10, 100), # 100 / 10 bytes per chunk = 10 chunks
+                          (200, 1, 100), # 100 / 200 bytes per chunk = 1 chunk
+                          ])
 def test_streaming_response_chunk_size(chunk_size, expected_chunks, data_size):
     test_data = b"0" * data_size
     content_chunks = []
 
     mock_response = Mock(spec=requests.Response)
+
     def mock_iter_content(chunk_size):
         # Simulate how requests would chunk the data.
         for i in range(0, len(test_data), chunk_size):
             chunk = test_data[i:i + chunk_size]
-            content_chunks.append(chunk)  # track chunks for verification
+            content_chunks.append(chunk) # track chunks for verification
             yield chunk
+
     mock_response.iter_content = mock_iter_content
     stream = _StreamingResponse(mock_response)
     stream.set_chunk_size(chunk_size)
-    
+
     # Read all data one byte at a time.
     received_data = b""
     while True:
-        chunk = stream.read(1) 
+        chunk = stream.read(1)
         if not chunk:
             break
         received_data += chunk
-    
-    assert received_data == test_data  # all data was received correctly
-    assert len(content_chunks) == expected_chunks  # correct number of chunks
-    assert all(len(c) <= chunk_size for c in content_chunks)  # chunks don't exceed size
+
+    assert received_data == test_data # all data was received correctly
+    assert len(content_chunks) == expected_chunks # correct number of chunks
+    assert all(len(c) <= chunk_size for c in content_chunks) # chunks don't exceed size