test fix

Author: jprakash-db

tests/unit/test_downloader.py

@@ -1,9 +1,11 @@
+from contextlib import contextmanager
 import unittest
 from unittest.mock import Mock, patch, MagicMock
 
 import requests
 
 import databricks.sql.cloudfetch.downloader as downloader
+from databricks.sql.common.http import DatabricksHttpClient
 from databricks.sql.exc import Error
 from databricks.sql.types import SSLOptions
 
@@ -12,6 +14,7 @@ def create_response(**kwargs) -> requests.Response:
     result = requests.Response()
     for k, v in kwargs.items():
         setattr(result, k, v)
+    result.close = Mock()
     return result
 
 
@@ -52,91 +55,94 @@ def test_run_link_past_expiry_buffer(self, mock_time):
 
         mock_time.assert_called_once()
 
-    @patch("requests.Session", return_value=MagicMock(get=MagicMock(return_value=None)))
     @patch("time.time", return_value=1000)
-    def test_run_get_response_not_ok(self, mock_time, mock_session):
-        mock_session.return_value.get.return_value = create_response(status_code=404)
-
+    def test_run_get_response_not_ok(self, mock_time):
+        http_client = DatabricksHttpClient.get_instance()
         settings = Mock(link_expiry_buffer_secs=0, download_timeout=0)
         settings.download_timeout = 0
         settings.use_proxy = False
         result_link = Mock(expiryTime=1001)
 
-        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))
+        with patch.object(
+            http_client,
+            "execute",
+            return_value=create_response(status_code=404, _content=b"1234"),
+        ):
+            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))
 
-    @patch("requests.Session", return_value=MagicMock(get=MagicMock(return_value=None)))
     @patch("time.time", return_value=1000)
-    def test_run_uncompressed_successful(self, mock_time, mock_session):
+    def test_run_uncompressed_successful(self, mock_time):
+        http_client = DatabricksHttpClient.get_instance()
         file_bytes = b"1234567890" * 10
-        mock_session.return_value.get.return_value = create_response(
-            status_code=200, _content=file_bytes
-        )
-
         settings = Mock(link_expiry_buffer_secs=0, download_timeout=0, use_proxy=False)
         settings.is_lz4_compressed = False
         result_link = Mock(bytesNum=100, expiryTime=1001)
 
-        d = downloader.ResultSetDownloadHandler(
-            settings, result_link, ssl_options=SSLOptions()
-        )
-        file = d.run()
+        with patch.object(
+            http_client,
+            "execute",
+            return_value=create_response(status_code=200, _content=file_bytes),
+        ):
+            d = downloader.ResultSetDownloadHandler(
+                settings, result_link, ssl_options=SSLOptions()
+            )
+            file = d.run()
 
-        assert file.file_bytes == b"1234567890" * 10
+            assert file.file_bytes == b"1234567890" * 10
 
-    @patch(
-        "requests.Session",
-        return_value=MagicMock(get=MagicMock(return_value=MagicMock(ok=True))),
-    )
     @patch("time.time", return_value=1000)
-    def test_run_compressed_successful(self, mock_time, mock_session):
+    def test_run_compressed_successful(self, mock_time):
+        http_client = DatabricksHttpClient.get_instance()
         file_bytes = b"1234567890" * 10
         compressed_bytes = b'\x04"M\x18h@d\x00\x00\x00\x00\x00\x00\x00#\x14\x00\x00\x00\xaf1234567890\n\x00BP67890\x00\x00\x00\x00'
-        mock_session.return_value.get.return_value = create_response(
-            status_code=200, _content=compressed_bytes
-        )
 
         settings = Mock(link_expiry_buffer_secs=0, download_timeout=0, use_proxy=False)
         settings.is_lz4_compressed = True
         result_link = Mock(bytesNum=100, expiryTime=1001)
+        with patch.object(
+            http_client,
+            "execute",
+            return_value=create_response(status_code=200, _content=compressed_bytes),
+        ):
+            d = downloader.ResultSetDownloadHandler(
+                settings, result_link, ssl_options=SSLOptions()
+            )
+            file = d.run()
+
+            assert file.file_bytes == b"1234567890" * 10
 
-        d = downloader.ResultSetDownloadHandler(
-            settings, result_link, ssl_options=SSLOptions()
-        )
-        file = d.run()
-
-        assert file.file_bytes == b"1234567890" * 10
-
-    @patch("requests.Session.get", side_effect=ConnectionError("foo"))
     @patch("time.time", return_value=1000)
-    def test_download_connection_error(self, mock_time, mock_session):
+    def test_download_connection_error(self, mock_time):
+
+        http_client = DatabricksHttpClient.get_instance()
         settings = Mock(
             link_expiry_buffer_secs=0, use_proxy=False, is_lz4_compressed=True
         )
         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()
-        )
-        with self.assertRaises(ConnectionError):
-            d.run()
+        with patch.object(http_client, "execute", side_effect=ConnectionError("foo")):
+            d = downloader.ResultSetDownloadHandler(
+                settings, result_link, ssl_options=SSLOptions()
+            )
+            with self.assertRaises(ConnectionError):
+                d.run()
 
-    @patch("requests.Session.get", side_effect=TimeoutError("foo"))
     @patch("time.time", return_value=1000)
-    def test_download_timeout(self, mock_time, mock_session):
+    def test_download_timeout(self, mock_time):
+        http_client = DatabricksHttpClient.get_instance()
         settings = Mock(
             link_expiry_buffer_secs=0, use_proxy=False, is_lz4_compressed=True
         )
         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()
-        )
-        with self.assertRaises(TimeoutError):
-            d.run()
+        with patch.object(http_client, "execute", side_effect=TimeoutError("foo")):
+            d = downloader.ResultSetDownloadHandler(
+                settings, result_link, ssl_options=SSLOptions()
+            )
+            with self.assertRaises(TimeoutError):
+                d.run()