correct unit tests for download management

Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com>

Author: varun-edachali-dbx

src/databricks/sql/cloudfetch/download_manager.py

@@ -91,15 +91,19 @@ def cancel_tasks_from_offset(self, start_row_offset: int):
         def to_cancel(link: TSparkArrowResultLink) -> bool:
             return link.startRowOffset < start_row_offset
 
-        tasks_to_cancel = [task for task in self._download_tasks if to_cancel(task.link)]
+        tasks_to_cancel = [
+            task for task in self._download_tasks if to_cancel(task.link)
+        ]
         for task in tasks_to_cancel:
             task.cancel()
         logger.info(
             f"ResultFileDownloadManager: cancelled {len(tasks_to_cancel)} tasks from offset {start_row_offset}"
         )
 
         # Remove cancelled tasks from the download queue
-        tasks_to_keep = [task for task in self._download_tasks if not to_cancel(task.link)]
+        tasks_to_keep = [
+            task for task in self._download_tasks if not to_cancel(task.link)
+        ]
         self._download_tasks = tasks_to_keep
 
         pending_links_to_keep = [

src/databricks/sql/cloudfetch/downloader.py

@@ -89,9 +89,7 @@ def run(self) -> DownloadedFile:
         )
 
         # Check if link is already expired or is expiring
-        self._validate_link(
-            self.link, self.settings.link_expiry_buffer_secs
-        )
+        self._validate_link(self.link, self.settings.link_expiry_buffer_secs)
 
         session = requests.Session()
         session.mount("http://", HTTPAdapter(max_retries=retryPolicy))

tests/unit/test_download_manager.py

@@ -14,11 +14,15 @@ class DownloadManagerTests(unittest.TestCase):
     def create_download_manager(
         self, links, max_download_threads=10, lz4_compressed=True
     ):
+        def expiry_callback(link: TSparkArrowResultLink):
+            return None
+
         return download_manager.ResultFileDownloadManager(
             links,
             max_download_threads,
             lz4_compressed,
             ssl_options=SSLOptions(),
+            expiry_callback=expiry_callback,
         )
 
     def create_result_link(

tests/unit/test_downloader.py

@@ -5,6 +5,7 @@
 
 import databricks.sql.cloudfetch.downloader as downloader
 from databricks.sql.exc import Error
+from databricks.sql.thrift_api.TCLIService.ttypes import TSparkArrowResultLink
 from databricks.sql.types import SSLOptions
 
 
@@ -20,15 +21,26 @@ class DownloaderTests(unittest.TestCase):
     Unit tests for checking downloader logic.
     """
 
+    def create_download_handler(
+        self, settings: Mock, result_link: Mock
+    ) -> downloader.ResultSetDownloadHandler:
+        def expiry_callback(link: TSparkArrowResultLink):
+            raise Error("Cloudfetch link has expired")
+
+        return downloader.ResultSetDownloadHandler(
+            settings,
+            result_link,
+            ssl_options=SSLOptions(),
+            expiry_callback=expiry_callback,
+        )
+
     @patch("time.time", return_value=1000)
     def test_run_link_expired(self, mock_time):
         settings = Mock()
         result_link = Mock()
         # Already expired
         result_link.expiryTime = 999
-        d = downloader.ResultSetDownloadHandler(
-            settings, result_link, ssl_options=SSLOptions()
-        )
+        d = self.create_download_handler(settings, result_link)
 
         with self.assertRaises(Error) as context:
             d.run()
@@ -42,9 +54,7 @@ def test_run_link_past_expiry_buffer(self, mock_time):
         result_link = Mock()
         # Within the expiry buffer time
         result_link.expiryTime = 1004
-        d = downloader.ResultSetDownloadHandler(
-            settings, result_link, ssl_options=SSLOptions()
-        )
+        d = self.create_download_handler(settings, result_link)
 
         with self.assertRaises(Error) as context:
             d.run()
@@ -62,9 +72,7 @@ def test_run_get_response_not_ok(self, mock_time, mock_session):
         settings.use_proxy = False
         result_link = Mock(expiryTime=1001)
 
-        d = downloader.ResultSetDownloadHandler(
-            settings, result_link, ssl_options=SSLOptions()
-        )
+        d = self.create_download_handler(settings, result_link)
         with self.assertRaises(requests.exceptions.HTTPError) as context:
             d.run()
         self.assertTrue("404" in str(context.exception))
@@ -81,9 +89,7 @@ def test_run_uncompressed_successful(self, mock_time, mock_session):
         settings.is_lz4_compressed = False
         result_link = Mock(bytesNum=100, expiryTime=1001)
 
-        d = downloader.ResultSetDownloadHandler(
-            settings, result_link, ssl_options=SSLOptions()
-        )
+        d = self.create_download_handler(settings, result_link)
         file = d.run()
 
         assert file.file_bytes == b"1234567890" * 10
@@ -104,9 +110,7 @@ def test_run_compressed_successful(self, mock_time, mock_session):
         settings.is_lz4_compressed = True
         result_link = Mock(bytesNum=100, expiryTime=1001)
 
-        d = downloader.ResultSetDownloadHandler(
-            settings, result_link, ssl_options=SSLOptions()
-        )
+        d = self.create_download_handler(settings, result_link)
         file = d.run()
 
         assert file.file_bytes == b"1234567890" * 10
@@ -120,9 +124,7 @@ def test_download_connection_error(self, mock_time, mock_session):
         result_link = Mock(bytesNum=100, expiryTime=1001)
         mock_session.return_value.get.return_value.content = b'\x04"M\x18h@d\x00\x00\x00\x00\x00\x00\x00#\x14\x00\x00\x00\xaf1234567890\n\x00BP67890\x00\x00\x00\x00'
 
-        d = downloader.ResultSetDownloadHandler(
-            settings, result_link, ssl_options=SSLOptions()
-        )
+        d = self.create_download_handler(settings, result_link)
         with self.assertRaises(ConnectionError):
             d.run()
 
@@ -135,8 +137,6 @@ def test_download_timeout(self, mock_time, mock_session):
         result_link = Mock(bytesNum=100, expiryTime=1001)
         mock_session.return_value.get.return_value.content = b'\x04"M\x18h@d\x00\x00\x00\x00\x00\x00\x00#\x14\x00\x00\x00\xaf1234567890\n\x00BP67890\x00\x00\x00\x00'
 
-        d = downloader.ResultSetDownloadHandler(
-            settings, result_link, ssl_options=SSLOptions()
-        )
+        d = self.create_download_handler(settings, result_link)
         with self.assertRaises(TimeoutError):
             d.run()