Update existing tests

Signed-off-by: Levko Kravets <[email protected]>

Author: kravets-levko

tests/unit/test_cloud_fetch_queue.py

@@ -1,10 +1,10 @@
 import pyarrow
 import unittest
 from unittest.mock import MagicMock, patch
-from ssl import create_default_context
 
 from databricks.sql.thrift_api.TCLIService.ttypes import TSparkArrowResultLink
 import databricks.sql.utils as utils
+from databricks.sql.types import SSLOptions
 
 class CloudFetchQueueSuite(unittest.TestCase):
 
@@ -51,7 +51,7 @@ def test_initializer_adds_links(self, mock_create_next_table):
             schema_bytes,
             result_links=result_links,
             max_download_threads=10,
-            ssl_context=create_default_context(),
+            ssl_options=SSLOptions(),
         )
 
         assert len(queue.download_manager._pending_links) == 10
@@ -65,7 +65,7 @@ def test_initializer_no_links_to_add(self):
             schema_bytes,
             result_links=result_links,
             max_download_threads=10,
-            ssl_context=create_default_context(),
+            ssl_options=SSLOptions(),
         )
 
         assert len(queue.download_manager._pending_links) == 0
@@ -78,7 +78,7 @@ def test_create_next_table_no_download(self, mock_get_next_downloaded_file):
             MagicMock(),
             result_links=[],
             max_download_threads=10,
-            ssl_context=create_default_context(),
+            ssl_options=SSLOptions(),
         )
 
         assert queue._create_next_table() is None
@@ -95,7 +95,7 @@ def test_initializer_create_next_table_success(self, mock_get_next_downloaded_fi
             result_links=[],
             description=description,
             max_download_threads=10,
-            ssl_context=create_default_context(),
+            ssl_options=SSLOptions(),
         )
         expected_result = self.make_arrow_table()
 
@@ -120,7 +120,7 @@ def test_next_n_rows_0_rows(self, mock_create_next_table):
             result_links=[],
             description=description,
             max_download_threads=10,
-            ssl_context=create_default_context(),
+            ssl_options=SSLOptions(),
         )
         assert queue.table == self.make_arrow_table()
         assert queue.table.num_rows == 4
@@ -140,7 +140,7 @@ def test_next_n_rows_partial_table(self, mock_create_next_table):
             result_links=[],
             description=description,
             max_download_threads=10,
-            ssl_context=create_default_context(),
+            ssl_options=SSLOptions(),
         )
         assert queue.table == self.make_arrow_table()
         assert queue.table.num_rows == 4
@@ -160,7 +160,7 @@ def test_next_n_rows_more_than_one_table(self, mock_create_next_table):
             result_links=[],
             description=description,
             max_download_threads=10,
-            ssl_context=create_default_context(),
+            ssl_options=SSLOptions(),
         )
         assert queue.table == self.make_arrow_table()
         assert queue.table.num_rows == 4
@@ -180,7 +180,7 @@ def test_next_n_rows_only_one_table_returned(self, mock_create_next_table):
             result_links=[],
             description=description,
             max_download_threads=10,
-            ssl_context=create_default_context(),
+            ssl_options=SSLOptions(),
         )
         assert queue.table == self.make_arrow_table()
         assert queue.table.num_rows == 4
@@ -199,7 +199,7 @@ def test_next_n_rows_empty_table(self, mock_create_next_table):
             result_links=[],
             description=description,
             max_download_threads=10,
-            ssl_context=create_default_context(),
+            ssl_options=SSLOptions(),
         )
         assert queue.table is None
 
@@ -216,7 +216,7 @@ def test_remaining_rows_empty_table_fully_returned(self, mock_create_next_table)
             result_links=[],
             description=description,
             max_download_threads=10,
-            ssl_context=create_default_context(),
+            ssl_options=SSLOptions(),
         )
         assert queue.table == self.make_arrow_table()
         assert queue.table.num_rows == 4
@@ -235,7 +235,7 @@ def test_remaining_rows_partial_table_fully_returned(self, mock_create_next_tabl
             result_links=[],
             description=description,
             max_download_threads=10,
-            ssl_context=create_default_context(),
+            ssl_options=SSLOptions(),
         )
         assert queue.table == self.make_arrow_table()
         assert queue.table.num_rows == 4
@@ -254,7 +254,7 @@ def test_remaining_rows_one_table_fully_returned(self, mock_create_next_table):
             result_links=[],
             description=description,
             max_download_threads=10,
-            ssl_context=create_default_context(),
+            ssl_options=SSLOptions(),
         )
         assert queue.table == self.make_arrow_table()
         assert queue.table.num_rows == 4
@@ -273,7 +273,7 @@ def test_remaining_rows_multiple_tables_fully_returned(self, mock_create_next_ta
             result_links=[],
             description=description,
             max_download_threads=10,
-            ssl_context=create_default_context(),
+            ssl_options=SSLOptions(),
         )
         assert queue.table == self.make_arrow_table()
         assert queue.table.num_rows == 4
@@ -293,7 +293,7 @@ def test_remaining_rows_empty_table(self, mock_create_next_table):
             result_links=[],
             description=description,
             max_download_threads=10,
-            ssl_context=create_default_context(),
+            ssl_options=SSLOptions(),
         )
         assert queue.table is None
 

tests/unit/test_download_manager.py

@@ -1,9 +1,8 @@
 import unittest
 from unittest.mock import patch, MagicMock
 
-from ssl import create_default_context
-
 import databricks.sql.cloudfetch.download_manager as download_manager
+from databricks.sql.types import SSLOptions
 from databricks.sql.thrift_api.TCLIService.ttypes import TSparkArrowResultLink
 
 
@@ -17,7 +16,7 @@ def create_download_manager(self, links, max_download_threads=10, lz4_compressed
             links,
             max_download_threads,
             lz4_compressed,
-            ssl_context=create_default_context(),
+            ssl_options=SSLOptions(),
         )
 
     def create_result_link(

tests/unit/test_downloader.py

@@ -2,10 +2,10 @@
 from unittest.mock import Mock, patch, MagicMock
 
 import requests
-from ssl import create_default_context
 
 import databricks.sql.cloudfetch.downloader as downloader
 from databricks.sql.exc import Error
+from databricks.sql.types import SSLOptions
 
 
 def create_response(**kwargs) -> requests.Response:
@@ -26,7 +26,7 @@ def test_run_link_expired(self, mock_time):
         result_link = Mock()
         # Already expired
         result_link.expiryTime = 999
-        d = downloader.ResultSetDownloadHandler(settings, result_link, ssl_context=create_default_context())
+        d = downloader.ResultSetDownloadHandler(settings, result_link, ssl_options=SSLOptions())
 
         with self.assertRaises(Error) as context:
             d.run()
@@ -40,7 +40,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_context=create_default_context())
+        d = downloader.ResultSetDownloadHandler(settings, result_link, ssl_options=SSLOptions())
 
         with self.assertRaises(Error) as context:
             d.run()
@@ -58,7 +58,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_context=create_default_context())
+        d = downloader.ResultSetDownloadHandler(settings, result_link, ssl_options=SSLOptions())
         with self.assertRaises(requests.exceptions.HTTPError) as context:
             d.run()
         self.assertTrue('404' in str(context.exception))
@@ -73,7 +73,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_context=create_default_context())
+        d = downloader.ResultSetDownloadHandler(settings, result_link, ssl_options=SSLOptions())
         file = d.run()
 
         assert file.file_bytes == b"1234567890" * 10
@@ -89,7 +89,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_context=create_default_context())
+        d = downloader.ResultSetDownloadHandler(settings, result_link, ssl_options=SSLOptions())
         file = d.run()
 
         assert file.file_bytes == b"1234567890" * 10
@@ -102,7 +102,7 @@ def test_download_connection_error(self, mock_time, mock_session):
         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_context=create_default_context())
+        d = downloader.ResultSetDownloadHandler(settings, result_link, ssl_options=SSLOptions())
         with self.assertRaises(ConnectionError):
             d.run()
 
@@ -114,6 +114,6 @@ def test_download_timeout(self, mock_time, mock_session):
         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_context=create_default_context())
+        d = downloader.ResultSetDownloadHandler(settings, result_link, ssl_options=SSLOptions())
         with self.assertRaises(TimeoutError):
             d.run()