simplifying tests

Signed-off-by: Sreekanth Vadigi <[email protected]>

Author: sreekanth-db

tests/unit/test_streaming_put.py

@@ -1,9 +1,7 @@
-
 import io
 from unittest.mock import patch, Mock, MagicMock
 
 import pytest
-import requests
 
 import databricks.sql.client as client
 
@@ -12,143 +10,104 @@ class TestStreamingPut:
     """Unit tests for streaming PUT functionality."""
 
     @pytest.fixture
-    def mock_connection(self):
-        return Mock()
-
-    @pytest.fixture
-    def mock_backend(self):
-        return Mock()
-
-    @pytest.fixture
-    def cursor(self, mock_connection, mock_backend):
-        return client.Cursor(
-            connection=mock_connection,
-            backend=mock_backend
-        )
+    def cursor(self):
+        return client.Cursor(connection=Mock(), backend=Mock())
 
     def _setup_mock_staging_put_stream_response(self, mock_backend):
         """Helper method to set up mock staging PUT stream response."""
         mock_result_set = Mock()
         mock_result_set.is_staging_operation = True
         mock_backend.execute_command.return_value = mock_result_set
-        
+
         mock_row = Mock()
         mock_row.operation = "PUT"
         mock_row.localFile = "__input_stream__"
         mock_row.presignedUrl = "https://example.com/upload"
         mock_row.headers = "{}"
         mock_result_set.fetchone.return_value = mock_row
-        
+
         return mock_result_set
 
-    def test_execute_with_valid_stream(self, cursor, mock_backend):
+    def test_execute_with_valid_stream(self, cursor):
         """Test execute method with valid input stream."""
-        
+
         # Mock the backend response
-        self._setup_mock_staging_put_stream_response(mock_backend)
-        
+        self._setup_mock_staging_put_stream_response(cursor.backend)
+
         # Test with valid stream
         test_stream = io.BytesIO(b"test data")
-        
-        with patch.object(cursor, '_handle_staging_put_stream') as mock_handler:
+
+        with patch.object(cursor, "_handle_staging_put_stream") as mock_handler:
             cursor.execute(
                 "PUT '__input_stream__' INTO '/Volumes/test/cat/schema/vol/file.txt'",
-                input_stream=test_stream
+                input_stream=test_stream,
             )
-            
+
             # Verify staging handler was called
             mock_handler.assert_called_once()
 
-    def test_execute_with_invalid_stream_types(self, cursor, mock_backend):
-
-        # Mock the backend response
-        self._setup_mock_staging_put_stream_response(mock_backend)
-        
-        # Test with None input stream
-        with pytest.raises(client.ProgrammingError) as excinfo:
-            cursor.execute(
-                "PUT '__input_stream__' INTO '/Volumes/test/cat/schema/vol/file.txt'",
-                input_stream=None
-            )
-        assert "No input stream provided for streaming operation" in str(excinfo.value)
-
-    def test_execute_with_none_stream_for_staging_put(self, cursor, mock_backend):
+    def test_execute_with_none_stream_for_staging_put(self, cursor):
         """Test execute method rejects None stream for streaming PUT operations."""
-        
+
         # Mock staging operation response for None case
-        self._setup_mock_staging_put_stream_response(mock_backend)
-        
+        self._setup_mock_staging_put_stream_response(cursor.backend)
+
         # None with __input_stream__ raises ProgrammingError
         with pytest.raises(client.ProgrammingError) as excinfo:
             cursor.execute(
                 "PUT '__input_stream__' INTO '/Volumes/test/cat/schema/vol/file.txt'",
-                input_stream=None
+                input_stream=None,
             )
         error_msg = str(excinfo.value)
         assert "No input stream provided for streaming operation" in error_msg
 
     def test_handle_staging_put_stream_success(self, cursor):
         """Test successful streaming PUT operation."""
-        
+
         presigned_url = "https://example.com/upload"
         headers = {"Content-Type": "text/plain"}
-        
-        with patch.object(cursor.connection.http_client, 'request') as mock_http_request:
+
+        with patch.object(
+            cursor.connection.http_client, "request"
+        ) as mock_http_request:
             mock_response = MagicMock()
             mock_response.status = 200
             mock_response.data = b"success"
             mock_http_request.return_value = mock_response
-            
+
             test_stream = io.BytesIO(b"test data")
             cursor._handle_staging_put_stream(
-                presigned_url=presigned_url,
-                stream=test_stream,
-                headers=headers
+                presigned_url=presigned_url, stream=test_stream, headers=headers
             )
-            
+
             # Verify the HTTP client was called correctly
             mock_http_request.assert_called_once()
             call_args = mock_http_request.call_args
             # Check positional arguments: (method, url, body=..., headers=...)
-            assert call_args[0][0].value == 'PUT'  # First positional arg is method
+            assert call_args[0][0].value == "PUT"  # First positional arg is method
             assert call_args[0][1] == presigned_url  # Second positional arg is url
             # Check keyword arguments
-            assert call_args[1]['body'] == b"test data"
-            assert call_args[1]['headers'] == headers
+            assert call_args[1]["body"] == b"test data"
+            assert call_args[1]["headers"] == headers
 
     def test_handle_staging_put_stream_http_error(self, cursor):
         """Test streaming PUT operation with HTTP error."""
-        
+
         presigned_url = "https://example.com/upload"
-        
-        with patch.object(cursor.connection.http_client, 'request') as mock_http_request:
+
+        with patch.object(
+            cursor.connection.http_client, "request"
+        ) as mock_http_request:
             mock_response = MagicMock()
             mock_response.status = 500
             mock_response.data = b"Internal Server Error"
             mock_http_request.return_value = mock_response
-            
+
             test_stream = io.BytesIO(b"test data")
             with pytest.raises(client.OperationalError) as excinfo:
                 cursor._handle_staging_put_stream(
-                    presigned_url=presigned_url,
-                    stream=test_stream
+                    presigned_url=presigned_url, stream=test_stream
                 )
-            
+
             # Check for the actual error message format
             assert "500" in str(excinfo.value)
-
-    def test_handle_staging_put_stream_network_error(self, cursor):
-        """Test streaming PUT operation with network error."""
-        
-        presigned_url = "https://example.com/upload"
-        
-        with patch.object(cursor.connection.http_client, 'request') as mock_http_request:
-            mock_http_request.side_effect = requests.exceptions.RequestException("Network error")
-            test_stream = io.BytesIO(b"test_data")
-            with pytest.raises(requests.exceptions.RequestException) as excinfo:
-                cursor._handle_staging_put_stream(
-                    presigned_url=presigned_url,
-                    stream=test_stream
-                )
-            
-            assert "Network error" in str(excinfo.value)