Fix test and formatting

Author: ksafonov-db

databricks/sdk/__init__.py

@@ -1,7 +1,7 @@
 import io
 import logging
-from abc import ABC, abstractmethod
 import urllib.parse
+from abc import ABC, abstractmethod
 from datetime import timedelta
 from types import TracebackType
 from typing import (Any, BinaryIO, Callable, Dict, Iterable, Iterator, List,
@@ -285,7 +285,9 @@ def _record_request_log(self, response: requests.Response, raw: bool = False) ->
             return
         logger.debug(RoundTrip(response, self._debug_headers, self._debug_truncate_bytes, raw).generate())
 
+
 class _RawResponse(ABC):
+
     @abstractmethod
     # follows Response signature: https://github.com/psf/requests/blob/main/src/requests/models.py#L799
     def iter_content(self, chunk_size: int = 1, decode_unicode: bool = False):
@@ -295,6 +297,7 @@ def iter_content(self, chunk_size: int = 1, decode_unicode: bool = False):
     def close(self):
         pass
 
+
 class _StreamingResponse(BinaryIO):
     _response: _RawResponse
     _buffer: bytes

databricks/sdk/_base_client.py

@@ -92,7 +92,8 @@ class Config:
     max_connections_per_pool: int = ConfigAttribute()
     databricks_environment: Optional[DatabricksEnvironment] = None
 
-    enable_experimental_files_api_client: bool = ConfigAttribute(env='DATABRICKS_ENABLE_EXPERIMENTAL_FILES_API_CLIENT')
+    enable_experimental_files_api_client: bool = ConfigAttribute(
+        env='DATABRICKS_ENABLE_EXPERIMENTAL_FILES_API_CLIENT')
     files_api_client_download_max_total_recovers = None
     files_api_client_download_max_total_recovers_without_progressing = 1
 

databricks/sdk/config.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import base64
+import logging
 import os
 import pathlib
 import platform
@@ -11,23 +12,25 @@
 from collections.abc import Iterator
 from io import BytesIO
 from types import TracebackType
-from typing import (TYPE_CHECKING, AnyStr, BinaryIO, Generator, Iterable, Optional, Type, Union)
+from typing import (TYPE_CHECKING, AnyStr, BinaryIO, Generator, Iterable,
+                    Optional, Type, Union)
 from urllib import parse
+
 from requests import RequestException
 
-import logging
+from .._base_client import _RawResponse, _StreamingResponse
 from .._property import _cached_property
 from ..errors import NotFound
 from ..service import files
 from ..service._internal import _escape_multi_segment_path_parameter
 from ..service.files import DownloadResponse
-from .._base_client import _RawResponse, _StreamingResponse
 
 if TYPE_CHECKING:
     from _typeshed import Self
 
 _LOG = logging.getLogger(__name__)
 
+
 class _DbfsIO(BinaryIO):
     MAX_CHUNK_SIZE = 1024 * 1024
 
@@ -699,13 +702,17 @@ def _download_raw_stream(self,
 
         result = DownloadResponse.from_dict(res)
         if not isinstance(result.contents, _StreamingResponse):
-            raise Exception("Internal error: response contents is of unexpected type: " + type(result.contents).__name__)
+            raise Exception("Internal error: response contents is of unexpected type: " +
+                            type(result.contents).__name__)
 
         return result
 
     def _wrap_stream(self, file_path: str, downloadResponse: DownloadResponse):
         underlying_response = _ResilientIterator._extract_raw_response(downloadResponse)
-        return _ResilientResponse(self, file_path, downloadResponse.last_modified, offset=0,
+        return _ResilientResponse(self,
+                                  file_path,
+                                  downloadResponse.last_modified,
+                                  offset=0,
                                   underlying_response=underlying_response)
 
 
@@ -728,8 +735,8 @@ def iter_content(self, chunk_size=1, decode_unicode=False):
             raise ValueError('Decode unicode is not supported')
 
         iterator = self.underlying_response.iter_content(chunk_size=chunk_size, decode_unicode=False)
-        self.iterator = _ResilientIterator(iterator, self.file_path, self.file_last_modified,
-                                           self.offset, self.api, chunk_size)
+        self.iterator = _ResilientIterator(iterator, self.file_path, self.file_last_modified, self.offset,
+                                           self.api, chunk_size)
         return self.iterator
 
     def close(self):
@@ -761,7 +768,6 @@ def __init__(self, underlying_iterator, file_path: str, file_last_modified: str,
         self._recovers_without_progressing_count: int = 0
         self._closed: bool = False
 
-
     def _should_recover(self) -> bool:
         if self._total_recovers_count == self._api._config.files_api_client_download_max_total_recovers:
             _LOG.debug("Total recovers limit exceeded")
@@ -773,7 +779,7 @@ def _should_recover(self) -> bool:
 
     def _recover(self) -> bool:
         if not self._should_recover():
-            return False  # recover suppressed, rethrow original exception
+            return False # recover suppressed, rethrow original exception
 
         self._total_recovers_count += 1
         self._recovers_without_progressing_count += 1
@@ -784,13 +790,15 @@ def _recover(self) -> bool:
             _LOG.debug("Trying to recover from offset " + str(self._offset))
 
             # following call includes all the required network retries
-            downloadResponse = self._api._download_raw_stream(self._file_path, self._offset, self._file_last_modified)
+            downloadResponse = self._api._download_raw_stream(self._file_path, self._offset,
+                                                              self._file_last_modified)
             underlying_response = _ResilientIterator._extract_raw_response(downloadResponse)
-            self._underlying_iterator = underlying_response.iter_content(chunk_size=self._chunk_size, decode_unicode=False)
+            self._underlying_iterator = underlying_response.iter_content(chunk_size=self._chunk_size,
+                                                                         decode_unicode=False)
             _LOG.debug("Recover succeeded")
             return True
         except:
-            return False  # recover failed, rethrow original exception
+            return False # recover failed, rethrow original exception
 
     def __next__(self):
         if self._closed:

databricks/sdk/mixins/files.py

@@ -8,14 +8,15 @@
 import requests
 
 from databricks.sdk import errors, useragent
-from databricks.sdk._base_client import _BaseClient, _StreamingResponse
+from databricks.sdk._base_client import (_BaseClient, _RawResponse,
+                                         _StreamingResponse)
 from databricks.sdk.core import DatabricksError
 
 from .clock import FakeClock
 from .fixture_server import http_fixture_server
 
 
-class DummyResponse(requests.Response):
+class DummyResponse(_RawResponse):
     _content: Iterator[bytes]
     _closed: bool = False
 
@@ -293,9 +294,9 @@ def test_streaming_response_chunk_size(chunk_size, expected_chunks, data_size):
     test_data = bytes(rng.getrandbits(8) for _ in range(data_size))
 
     content_chunks = []
-    mock_response = Mock(spec=requests.Response)
+    mock_response = Mock(spec=_RawResponse)
 
-    def mock_iter_content(chunk_size):
+    def mock_iter_content(chunk_size: int, decode_unicode: bool):
         # Simulate how requests would chunk the data.
         for i in range(0, len(test_data), chunk_size):
             chunk = test_data[i:i + chunk_size]