added mocks to resolve the errors caused by log_latency decorator in tests

Signed-off-by: Sai Shree Pradhan <saishree.pradhan@databricks.com>

Author: saishreeeee

tests/unit/test_client.py

@@ -8,19 +8,6 @@
 from datetime import datetime, date
 from uuid import UUID
 
-def noop_log_latency_decorator(*args, **kwargs):
-    """
-    This is a no-op decorator. It is used to patch the log_latency decorator
-    during tests, so that the tests for client logic are not affected by the
-    telemetry logging logic. It accepts any arguments and returns a decorator
-    that returns the original function unmodified.
-    """
-    def decorator(func):
-        return func
-    return decorator
-
-patch('databricks.sql.telemetry.latency_logger.log_latency', new=noop_log_latency_decorator).start()
-
 from databricks.sql.thrift_api.TCLIService.ttypes import (
     TOpenSessionResp,
     TExecuteStatementResp,
@@ -51,6 +38,10 @@ 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,
@@ -331,7 +322,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(), Mock())
+        cursor = client.Cursor(Mock(), ThriftBackendMockFactory.new())
         cursor.close()
 
         with self.assertRaises(Error) as e:
@@ -343,14 +334,19 @@ def test_closed_cursor_doesnt_allow_operations(self):
             self.assertIn("closed", e.msg)
 
     def test_negative_fetch_throws_exception(self):
-        result_set = client.ResultSet(Mock(), Mock(), Mock())
+        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())
 
         with self.assertRaises(ValueError) as e:
             result_set.fetchmany(-1)
 
     def test_context_manager_closes_cursor(self):
         mock_close = Mock()
-        with client.Cursor(Mock(), Mock()) as cursor:
+        with client.Cursor(Mock(), ThriftBackendMockFactory.new()) as cursor:
             cursor.close = mock_close
         mock_close.assert_called_once_with()
 
@@ -393,7 +389,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 = Mock()
+                mock_thrift_backend = ThriftBackendMockFactory.new()
 
                 cursor = client.Cursor(Mock(), mock_thrift_backend)
                 cursor.schemas(**req_args)
@@ -416,7 +412,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 = Mock()
+                mock_thrift_backend = ThriftBackendMockFactory.new()
 
                 cursor = client.Cursor(Mock(), mock_thrift_backend)
                 cursor.tables(**req_args)
@@ -439,7 +435,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 = Mock()
+                mock_thrift_backend = ThriftBackendMockFactory.new()
 
                 cursor = client.Cursor(Mock(), mock_thrift_backend)
                 cursor.columns(**req_args)

tests/unit/test_fetches.py

@@ -7,19 +7,6 @@
 except ImportError:
     pa = None
 
-def noop_log_latency_decorator(*args, **kwargs):
-    """
-    This is a no-op decorator. It is used to patch the log_latency decorator
-    during tests, so that the tests for fetches logic are not affected by the
-    telemetry logging logic. It accepts any arguments and returns a decorator
-    that returns the original function unmodified.
-    """
-    def decorator(func):
-        return func
-    return decorator
-
-patch('databricks.sql.telemetry.latency_logger.log_latency', new=noop_log_latency_decorator).start()
-
 import databricks.sql.client as client
 from databricks.sql.utils import ExecuteResponse, ArrowQueue
 
@@ -71,6 +58,14 @@ def make_dummy_result_set_from_initial_results(initial_results):
             for col_id in range(num_cols)
         ]
         return rs
+    
+    @staticmethod
+    def mock_thrift_backend_with_retry_policy():
+        mock_thrift_backend = Mock()
+        mock_retry_policy = Mock()
+        mock_retry_policy.history = []
+        mock_thrift_backend.retry_policy = mock_retry_policy
+        return mock_thrift_backend
 
     @staticmethod
     def make_dummy_result_set_from_batch_list(batch_list):
@@ -92,7 +87,7 @@ def fetch_results(
 
             return results, batch_index < len(batch_list)
 
-        mock_thrift_backend = Mock()
+        mock_thrift_backend = FetchTests.mock_thrift_backend_with_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