removed the changes in test_client and test_fetches

saishreeeee · saishreeeee · commit 2a26965d71b4 · 2025-07-02T15:49:14.000+05:30
Signed-off-by: Sai Shree Pradhan &lt;saishree.pradhan@databricks.com&gt;
diff --git a/tests/unit/test_client.py b/tests/unit/test_client.py
@@ -38,10 +38,6 @@ def new(cls):
         cls.apply_property_to_mock(ThriftBackendMock, staging_allowed_local_path=None)
         MockTExecuteStatementResp = MagicMock(spec=TExecuteStatementResp())
 
-        mock_retry_policy = Mock()
-        mock_retry_policy.history = []
-        cls.apply_property_to_mock(ThriftBackendMock, retry_policy=mock_retry_policy)
-
         cls.apply_property_to_mock(
             MockTExecuteStatementResp,
             description=None,
@@ -322,7 +318,7 @@ def test_executing_multiple_commands_uses_the_most_recent_command(
         mock_result_sets[1].fetchall.assert_called_once_with()
 
     def test_closed_cursor_doesnt_allow_operations(self):
-        cursor = client.Cursor(Mock(), ThriftBackendMockFactory.new())
+        cursor = client.Cursor(Mock(), Mock())
         cursor.close()
 
         with self.assertRaises(Error) as e:
@@ -334,19 +330,14 @@ def test_closed_cursor_doesnt_allow_operations(self):
             self.assertIn("closed", e.msg)
 
     def test_negative_fetch_throws_exception(self):
-        mock_connection = Mock()
-        mock_connection.get_session_id_hex.return_value = "test_session"
-        mock_execute_response = Mock()
-        mock_execute_response.command_handle = None  
-        
-        result_set = client.ResultSet(mock_connection, mock_execute_response, ThriftBackendMockFactory.new())
+        result_set = client.ResultSet(Mock(), Mock(), Mock())
 
         with self.assertRaises(ValueError) as e:
             result_set.fetchmany(-1)
 
     def test_context_manager_closes_cursor(self):
         mock_close = Mock()
-        with client.Cursor(Mock(), ThriftBackendMockFactory.new()) as cursor:
+        with client.Cursor(Mock(), Mock()) as cursor:
             cursor.close = mock_close
         mock_close.assert_called_once_with()
 
@@ -389,7 +380,7 @@ def test_get_schemas_parameters_passed_to_thrift_backend(self, mock_thrift_backe
         for req_args in req_args_combinations:
             req_args = {k: v for k, v in req_args.items() if v != "NOT_SET"}
             with self.subTest(req_args=req_args):
-                mock_thrift_backend = ThriftBackendMockFactory.new()
+                mock_thrift_backend = Mock()
 
                 cursor = client.Cursor(Mock(), mock_thrift_backend)
                 cursor.schemas(**req_args)
@@ -412,7 +403,7 @@ def test_get_tables_parameters_passed_to_thrift_backend(self, mock_thrift_backen
         for req_args in req_args_combinations:
             req_args = {k: v for k, v in req_args.items() if v != "NOT_SET"}
             with self.subTest(req_args=req_args):
-                mock_thrift_backend = ThriftBackendMockFactory.new()
+                mock_thrift_backend = Mock()
 
                 cursor = client.Cursor(Mock(), mock_thrift_backend)
                 cursor.tables(**req_args)
@@ -435,7 +426,7 @@ def test_get_columns_parameters_passed_to_thrift_backend(self, mock_thrift_backe
         for req_args in req_args_combinations:
             req_args = {k: v for k, v in req_args.items() if v != "NOT_SET"}
             with self.subTest(req_args=req_args):
-                mock_thrift_backend = ThriftBackendMockFactory.new()
+                mock_thrift_backend = Mock()
 
                 cursor = client.Cursor(Mock(), mock_thrift_backend)
                 cursor.columns(**req_args)
diff --git a/tests/unit/test_fetches.py b/tests/unit/test_fetches.py
@@ -80,9 +80,6 @@ def fetch_results(
             return results, batch_index < len(batch_list)
 
         mock_thrift_backend = Mock()
-        mock_retry_policy = Mock()
-        mock_retry_policy.history = []
-        mock_thrift_backend.retry_policy = mock_retry_policy
         mock_thrift_backend.fetch_results = fetch_results
         num_cols = len(batch_list[0][0]) if batch_list and batch_list[0] else 0
 
