fixed (most) tests by accounting for normalised Session interface

varun-edachali-dbx · varun-edachali-dbx · commit 7c733eeb0fcd · 2025-05-28T03:50:50.000Z
have not considered how to fix the protocol version check yet

Signed-off-by: varun-edachali-dbx &lt;varun.edachali@databricks.com&gt;
diff --git a/tests/e2e/test_driver.py b/tests/e2e/test_driver.py
@@ -824,9 +824,7 @@ def test_close_connection_closes_cursors(self):
             status_request = ttypes.TGetOperationStatusReq(
                 operationHandle=ars.command_id, getProgressUpdate=False
             )
-            op_status_at_server = ars.backend._client.GetOperationStatus(
-                status_request
-            )
+            op_status_at_server = ars.backend._client.GetOperationStatus(status_request)
             assert (
                 op_status_at_server.operationState
                 != ttypes.TOperationState.CLOSED_STATE
@@ -856,17 +854,19 @@ def test_closing_a_closed_connection_doesnt_fail(self, caplog):
                     raise KeyboardInterrupt("Simulated interrupt")
         finally:
             if conn is not None:
-                assert not conn.open, "Connection should be closed after KeyboardInterrupt"
+                assert (
+                    not conn.open
+                ), "Connection should be closed after KeyboardInterrupt"
 
     def test_cursor_close_properly_closes_operation(self):
         """Test that Cursor.close() properly closes the active operation handle on the server."""
         with self.connection() as conn:
             cursor = conn.cursor()
             try:
                 cursor.execute("SELECT 1 AS test")
-                assert cursor.active_op_handle is not None
+                assert cursor.active_command_id is not None
                 cursor.close()
-                assert cursor.active_op_handle is None
+                assert cursor.active_command_id is None
                 assert not cursor.open
             finally:
                 if cursor.open:
@@ -883,26 +883,28 @@ def test_cursor_close_properly_closes_operation(self):
                         raise KeyboardInterrupt("Simulated interrupt")
         finally:
             if cursor is not None:
-                assert not cursor.open, "Cursor should be closed after KeyboardInterrupt"
+                assert (
+                    not cursor.open
+                ), "Cursor should be closed after KeyboardInterrupt"
 
     def test_nested_cursor_context_managers(self):
         """Test that nested cursor context managers properly close operations on the server."""
         with self.connection() as conn:
             with conn.cursor() as cursor1:
                 cursor1.execute("SELECT 1 AS test1")
-                assert cursor1.active_op_handle is not None
+                assert cursor1.active_command_id is not None
 
                 with conn.cursor() as cursor2:
                     cursor2.execute("SELECT 2 AS test2")
-                    assert cursor2.active_op_handle is not None
+                    assert cursor2.active_command_id is not None
 
                 # After inner context manager exit, cursor2 should be not open
                 assert not cursor2.open
-                assert cursor2.active_op_handle is None
+                assert cursor2.active_command_id is None
 
             # After outer context manager exit, cursor1 should be not open
             assert not cursor1.open
-            assert cursor1.active_op_handle is None
+            assert cursor1.active_command_id is None
 
     def test_cursor_error_handling(self):
         """Test that cursor close handles errors properly to prevent orphaned operations."""
@@ -911,7 +913,7 @@ def test_cursor_error_handling(self):
 
             cursor.execute("SELECT 1 AS test")
 
-            op_handle = cursor.active_op_handle
+            op_handle = cursor.active_command_id
 
             assert op_handle is not None
 
diff --git a/tests/unit/test_client.py b/tests/unit/test_client.py
@@ -81,7 +81,10 @@ class ClientTestSuite(unittest.TestCase):
         "access_token": "tok",
     }
 
-    @patch("%s.session.ThriftDatabricksClient" % PACKAGE_NAME, ThriftDatabricksClientMockFactory.new())
+    @patch(
+        "%s.session.ThriftDatabricksClient" % PACKAGE_NAME,
+        ThriftDatabricksClientMockFactory.new(),
+    )
     @patch("%s.client.ResultSet" % PACKAGE_NAME)
     def test_closing_connection_closes_commands(self, mock_result_set_class):
         # Test once with has_been_closed_server side, once without
@@ -294,10 +297,10 @@ def test_get_columns_parameters_passed_to_thrift_backend(self, mock_thrift_backe
     def test_cancel_command_calls_the_backend(self):
         mock_thrift_backend = Mock()
         cursor = client.Cursor(Mock(), mock_thrift_backend)
-        mock_op_handle = Mock()
-        cursor.active_op_handle = mock_op_handle
+        mock_command_id = Mock()
+        cursor.active_command_id = mock_command_id
         cursor.cancel()
-        mock_thrift_backend.cancel_command.assert_called_with(mock_op_handle)
+        mock_thrift_backend.cancel_command.assert_called_with(mock_command_id)
 
     @patch("databricks.sql.client.logger")
     def test_cancel_command_will_issue_warning_for_cancel_with_no_executing_command(
@@ -531,9 +534,13 @@ def test_access_current_query_id(self):
 
         self.assertIsNone(cursor.query_id)
 
-        cursor.active_op_handle = TOperationHandle(
-            operationId=THandleIdentifier(guid=UUID(operation_id).bytes, secret=0x00),
-            operationType=TOperationType.EXECUTE_STATEMENT,
+        cursor.active_command_id = CommandId.from_thrift_handle(
+            TOperationHandle(
+                operationId=THandleIdentifier(
+                    guid=UUID(operation_id).bytes, secret=0x00
+                ),
+                operationType=TOperationType.EXECUTE_STATEMENT,
+            )
         )
         self.assertEqual(cursor.query_id.upper(), operation_id.upper())
 
@@ -544,70 +551,72 @@ def test_cursor_close_handles_exception(self):
         """Test that Cursor.close() handles exceptions from close_command properly."""
         mock_backend = Mock()
         mock_connection = Mock()
-        mock_op_handle = Mock()
-        
+        mock_command_id = Mock()
+
         mock_backend.close_command.side_effect = Exception("Test error")
 
         cursor = client.Cursor(mock_connection, mock_backend)
-        cursor.active_op_handle = mock_op_handle
+        cursor.active_command_id = mock_command_id
 
         cursor.close()
 
-        mock_backend.close_command.assert_called_once_with(mock_op_handle)
-        
-        self.assertIsNone(cursor.active_op_handle)
-        
+        mock_backend.close_command.assert_called_once_with(mock_command_id)
+
+        self.assertIsNone(cursor.active_command_id)
+
         self.assertFalse(cursor.open)
 
     def test_cursor_context_manager_handles_exit_exception(self):
         """Test that cursor's context manager handles exceptions during __exit__."""
         mock_backend = Mock()
         mock_connection = Mock()
-        
+
         cursor = client.Cursor(mock_connection, mock_backend)
         original_close = cursor.close
         cursor.close = Mock(side_effect=Exception("Test error during close"))
-        
+
         try:
             with cursor:
                 raise ValueError("Test error inside context")
         except ValueError:
             pass
-        
+
         cursor.close.assert_called_once()
 
     def test_connection_close_handles_cursor_close_exception(self):
         """Test that _close handles exceptions from cursor.close() properly."""
         cursors_closed = []
-        
+
         def mock_close_with_exception():
             cursors_closed.append(1)
             raise Exception("Test error during close")
-        
+
         cursor1 = Mock()
         cursor1.close = mock_close_with_exception
-        
+
         def mock_close_normal():
             cursors_closed.append(2)
-        
+
         cursor2 = Mock()
         cursor2.close = mock_close_normal
-        
+
         mock_backend = Mock()
         mock_session_handle = Mock()
-        
+
         try:
             for cursor in [cursor1, cursor2]:
                 try:
                     cursor.close()
                 except Exception:
                     pass
-                    
+
             mock_backend.close_session(mock_session_handle)
         except Exception as e:
             self.fail(f"Connection close should handle exceptions: {e}")
-        
-        self.assertEqual(cursors_closed, [1, 2], "Both cursors should have close called")
+
+        self.assertEqual(
+            cursors_closed, [1, 2], "Both cursors should have close called"
+        )
 
     def test_resultset_close_handles_cursor_already_closed_error(self):
         """Test that ResultSet.close() handles CursorAlreadyClosedError properly."""
@@ -616,7 +625,7 @@ def test_resultset_close_handles_cursor_already_closed_error(self):
         result_set.backend.CLOSED_OP_STATE = 'CLOSED'
         result_set.connection = Mock()
         result_set.connection.open = True
-        result_set.op_state = 'RUNNING'
+        result_set.op_state = "RUNNING"
         result_set.has_been_closed_server_side = False
         result_set.command_id = Mock()
 
diff --git a/tests/unit/test_fetches.py b/tests/unit/test_fetches.py
@@ -66,7 +66,7 @@ def make_dummy_result_set_from_batch_list(batch_list):
         batch_index = 0
 
         def fetch_results(
-            op_handle,
+            command_id,
             max_rows,
             max_bytes,
             expected_row_start_offset,
diff --git a/tests/unit/test_session.py b/tests/unit/test_session.py
@@ -4,7 +4,10 @@
 
 from databricks.sql.thrift_api.TCLIService.ttypes import (
     TOpenSessionResp,
+    TSessionHandle,
+    THandleIdentifier,
 )
+from databricks.sql.ids import SessionId, BackendType
 
 import databricks.sql
 
@@ -25,16 +28,17 @@ class SessionTestSuite(unittest.TestCase):
     def test_close_uses_the_correct_session_id(self, mock_client_class):
         instance = mock_client_class.return_value
 
-        mock_open_session_resp = MagicMock(spec=TOpenSessionResp)()
-        mock_open_session_resp.sessionHandle.sessionId = b"\x22"
-        instance.open_session.return_value = mock_open_session_resp
+        # Create a mock SessionId that will be returned by open_session
+        mock_session_id = SessionId(BackendType.THRIFT, b"\x22", b"\x33")
+        instance.open_session.return_value = mock_session_id
 
         connection = databricks.sql.connect(**self.DUMMY_CONNECTION_ARGS)
         connection.close()
 
-        # Check the close session request has an id of x22
-        close_session_id = instance.close_session.call_args[0][0].sessionId
-        self.assertEqual(close_session_id, b"\x22")
+        # Check that close_session was called with the correct SessionId
+        close_session_call_args = instance.close_session.call_args[0][0]
+        self.assertEqual(close_session_call_args.guid, b"\x22")
+        self.assertEqual(close_session_call_args.secret, b"\x33")
 
     @patch("%s.session.ThriftDatabricksClient" % PACKAGE_NAME)
     def test_auth_args(self, mock_client_class):
@@ -112,16 +116,17 @@ def test_useragent_header(self, mock_client_class):
     def test_context_manager_closes_connection(self, mock_client_class):
         instance = mock_client_class.return_value
 
-        mock_open_session_resp = MagicMock(spec=TOpenSessionResp)()
-        mock_open_session_resp.sessionHandle.sessionId = b"\x22"
-        instance.open_session.return_value = mock_open_session_resp
+        # Create a mock SessionId that will be returned by open_session
+        mock_session_id = SessionId(BackendType.THRIFT, b"\x22", b"\x33")
+        instance.open_session.return_value = mock_session_id
 
         with databricks.sql.connect(**self.DUMMY_CONNECTION_ARGS) as connection:
             pass
 
-        # Check the close session request has an id of x22
-        close_session_id = instance.close_session.call_args[0][0].sessionId
-        self.assertEqual(close_session_id, b"\x22")
+        # Check that close_session was called with the correct SessionId
+        close_session_call_args = instance.close_session.call_args[0][0]
+        self.assertEqual(close_session_call_args.guid, b"\x22")
+        self.assertEqual(close_session_call_args.secret, b"\x33")
 
     @patch("%s.session.ThriftDatabricksClient" % PACKAGE_NAME)
     def test_max_number_of_retries_passthrough(self, mock_client_class):
@@ -141,46 +146,54 @@ def test_socket_timeout_passthrough(self, mock_client_class):
     @patch("%s.session.ThriftDatabricksClient" % PACKAGE_NAME)
     def test_configuration_passthrough(self, mock_client_class):
         mock_session_config = Mock()
+
+        # Create a mock SessionId that will be returned by open_session
+        mock_session_id = SessionId(BackendType.THRIFT, b"\x22", b"\x33")
+        mock_client_class.return_value.open_session.return_value = mock_session_id
+
         databricks.sql.connect(
             session_configuration=mock_session_config, **self.DUMMY_CONNECTION_ARGS
         )
 
-        self.assertEqual(
-            mock_client_class.return_value.open_session.call_args[0][0],
-            mock_session_config,
-        )
+        # Check that open_session was called with the correct session_configuration
+        call_kwargs = mock_client_class.return_value.open_session.call_args[1]
+        self.assertEqual(call_kwargs["session_configuration"], mock_session_config)
 
     @patch("%s.session.ThriftDatabricksClient" % PACKAGE_NAME)
     def test_initial_namespace_passthrough(self, mock_client_class):
         mock_cat = Mock()
         mock_schem = Mock()
 
+        # Create a mock SessionId that will be returned by open_session
+        mock_session_id = SessionId(BackendType.THRIFT, b"\x22", b"\x33")
+        mock_client_class.return_value.open_session.return_value = mock_session_id
+
         databricks.sql.connect(
             **self.DUMMY_CONNECTION_ARGS, catalog=mock_cat, schema=mock_schem
         )
-        self.assertEqual(
-            mock_client_class.return_value.open_session.call_args[0][1], mock_cat
-        )
-        self.assertEqual(
-            mock_client_class.return_value.open_session.call_args[0][2], mock_schem
-        )
+
+        # Check that open_session was called with the correct catalog and schema
+        call_kwargs = mock_client_class.return_value.open_session.call_args[1]
+        self.assertEqual(call_kwargs["catalog"], mock_cat)
+        self.assertEqual(call_kwargs["schema"], mock_schem)
 
     @patch("%s.session.ThriftDatabricksClient" % PACKAGE_NAME)
     def test_finalizer_closes_abandoned_connection(self, mock_client_class):
         instance = mock_client_class.return_value
 
-        mock_open_session_resp = MagicMock(spec=TOpenSessionResp)()
-        mock_open_session_resp.sessionHandle.sessionId = b"\x22"
-        instance.open_session.return_value = mock_open_session_resp
+        # Create a mock SessionId that will be returned by open_session
+        mock_session_id = SessionId(BackendType.THRIFT, b"\x22", b"\x33")
+        instance.open_session.return_value = mock_session_id
 
         databricks.sql.connect(**self.DUMMY_CONNECTION_ARGS)
 
         # not strictly necessary as the refcount is 0, but just to be sure
         gc.collect()
 
-        # Check the close session request has an id of x22
-        close_session_id = instance.close_session.call_args[0][0].sessionId
-        self.assertEqual(close_session_id, b"\x22")
+        # Check that close_session was called with the correct SessionId
+        close_session_call_args = instance.close_session.call_args[0][0]
+        self.assertEqual(close_session_call_args.guid, b"\x22")
+        self.assertEqual(close_session_call_args.secret, b"\x33")
 
 
 if __name__ == "__main__":
diff --git a/tests/unit/test_thrift_backend.py b/tests/unit/test_thrift_backend.py
