@@ -89,32 +89,40 @@ class ClientTestSuite(unittest.TestCase):
8989
9090@patch ("%s.session.ThriftDatabricksClient" % PACKAGE_NAME )
9191def test_closing_connection_closes_commands (self ,mock_thrift_client_class ):
92- """Test that connection.close() properly closes result sets through the real close chain."""
93- # Test once with has_been_closed_server side, once without
92+ """Test that closing a connection properly closes commands.
93+
94+ This test verifies that when a connection is closed:
95+ 1. the active result set is marked as closed server-side
96+ 2. The operation state is set to CLOSED
97+ 3. backend.close_command is called only for commands that weren't already closed
98+
99+ Args:
100+ mock_thrift_client_class: Mock for ThriftBackend class
101+ """
102+
94103for closed in (True ,False ):
95104with self .subTest (closed = closed ):
105+ # set initial state based on whether the command is already closed
106+ initial_state = (
107+ CommandState .CLOSED if closed else CommandState .SUCCEEDED
108+ )
109+
96110# Mock the execute response with controlled state
97111mock_execute_response = Mock (spec = ExecuteResponse )
98-
99- mock_execute_response .command_id = Mock (spec = CommandId )
100- mock_execute_response .status = (
101- CommandState .SUCCEEDED if not closed else CommandState .CLOSED
102- )
112+ mock_execute_response .status = initial_state
103113mock_execute_response .has_been_closed_server_side = closed
104114mock_execute_response .is_staging_operation = False
115+ mock_execute_response .command_id = Mock (spec = CommandId )
105116
106117# Mock the backend that will be used by the real ThriftResultSet
107118mock_backend = Mock (spec = ThriftDatabricksClient )
108119mock_backend .staging_allowed_local_path = None
109-
110- # Configure the decorator's mock to return our specific mock_backend
111120mock_thrift_client_class .return_value = mock_backend
112121
113122# Create connection and cursor
114123connection = databricks .sql .connect (** self .DUMMY_CONNECTION_ARGS )
115124cursor = connection .cursor ()
116125
117- # Create a REAL ThriftResultSet that will be returned by execute_command
118126real_result_set = ThriftResultSet (
119127connection = connection ,
120128execute_response = mock_execute_response ,
@@ -127,8 +135,7 @@ def test_closing_connection_closes_commands(self, mock_thrift_client_class):
127135# Execute a command - this should set cursor.active_result_set to our real result set
128136cursor .execute ("SELECT 1" )
129137
130- # Close the connection - this should trigger the real close chain:
131- # connection.close() -> cursor.close() -> result_set.close()
138+ # Close the connection
132139connection .close ()
133140
134141# Verify the REAL close logic worked through the chain: