Skip to content

Commit 54f2b93

Browse files
varun-edachali-dbxvarun-edachali-dbx
committed
simplify corrected test
Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com>
1 parent fde4634 commit 54f2b93

File tree

1 file changed

+54
-67
lines changed

1 file changed

+54
-67
lines changed

‎tests/unit/test_client.py‎

Lines changed: 54 additions & 67 deletions
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,9 @@
1414
TOperationHandle,
1515
THandleIdentifier,
1616
TOperationType,
17+
TOperationState,
1718
)
19+
from databricks.sql.thrift_api.TCLIService import ttypes
1820
from databricks.sql.backend.thrift_backend import ThriftDatabricksClient
1921

2022
import databricks.sql
@@ -83,32 +85,20 @@ class ClientTestSuite(unittest.TestCase):
8385
"access_token": "tok",
8486
}
8587

86-
@patch(
87-
"%s.backend.thrift_backend.ThriftDatabricksClient.close_command" % PACKAGE_NAME
88-
)
89-
def test_closing_connection_closes_commands(self, mock_close_command):
88+
@patch("%s.session.ThriftDatabricksClient" % PACKAGE_NAME)
89+
def test_closing_connection_closes_commands(self, mock_thrift_client_class):
9090
"""Test that connection.close() properly closes result sets through the real close chain."""
9191
# Test once with has_been_closed_server side, once without
9292
for closed in (True, False):
9393
with self.subTest(closed=closed):
94-
mock_close_command.reset_mock() # Reset for each subtest
95-
96-
# Create a real ThriftResultSet with mocked dependencies
97-
from databricks.sql.utils import ExecuteResponse
98-
from databricks.sql.backend.types import CommandId, CommandState
99-
from databricks.sql.result_set import ThriftResultSet
100-
10194
# Mock the execute response with controlled state
10295
mock_execute_response = Mock(spec=ExecuteResponse)
103-
mock_execute_response.command_id = Mock(spec=CommandId)
104-
105-
# Use actual Thrift operation states, not CommandState enums
106-
from databricks.sql.thrift_api.TCLIService import ttypes
10796

97+
mock_execute_response.command_id = Mock(spec=CommandId)
10898
mock_execute_response.status = (
109-
ttypes.TOperationState.FINISHED_STATE
99+
TOperationState.FINISHED_STATE
110100
if not closed
111-
else ttypes.TOperationState.CLOSED_STATE
101+
else TOperationState.CLOSED_STATE
112102
)
113103
mock_execute_response.has_been_closed_server_side = closed
114104
mock_execute_response.is_staging_operation = False
@@ -117,62 +107,59 @@ def test_closing_connection_closes_commands(self, mock_close_command):
117107
mock_backend = Mock(spec=ThriftDatabricksClient)
118108
mock_backend.staging_allowed_local_path = None
119109

110+
# Configure the decorator's mock to return our specific mock_backend
111+
mock_thrift_client_class.return_value = mock_backend
112+
120113
# Create connection and cursor
121-
with patch(
122-
f"{self.PACKAGE_NAME}.session.ThriftDatabricksClient",
123-
return_value=mock_backend,
124-
):
125-
connection = databricks.sql.connect(**self.DUMMY_CONNECTION_ARGS)
126-
cursor = connection.cursor()
127-
128-
# Create a REAL ThriftResultSet that will be returned by execute_command
129-
real_result_set = ThriftResultSet(
130-
connection=connection,
131-
execute_response=mock_execute_response,
132-
thrift_client=mock_backend,
133-
)
114+
connection = databricks.sql.connect(**self.DUMMY_CONNECTION_ARGS)
115+
cursor = connection.cursor()
116+
117+
# Create a REAL ThriftResultSet that will be returned by execute_command
118+
real_result_set = ThriftResultSet(
119+
connection=connection,
120+
execute_response=mock_execute_response,
121+
thrift_client=mock_backend,
122+
)
134123

135-
# Verify initial state
136-
self.assertEqual(
137-
real_result_set.has_been_closed_server_side, closed
138-
)
139-
expected_op_state = (
140-
CommandState.CLOSED if closed else CommandState.SUCCEEDED
141-
)
142-
self.assertEqual(real_result_set.op_state, expected_op_state)
124+
# Verify initial state
125+
self.assertEqual(real_result_set.has_been_closed_server_side, closed)
126+
expected_op_state = (
127+
CommandState.CLOSED if closed else CommandState.SUCCEEDED
128+
)
129+
self.assertEqual(real_result_set.op_state, expected_op_state)
143130

144-
# Mock execute_command to return our real result set
145-
cursor.backend.execute_command = Mock(return_value=real_result_set)
131+
# Mock execute_command to return our real result set
132+
cursor.backend.execute_command = Mock(return_value=real_result_set)
146133

147-
# Execute a command - this should set cursor.active_result_set to our real result set
148-
cursor.execute("SELECT 1")
134+
# Execute a command - this should set cursor.active_result_set to our real result set
135+
cursor.execute("SELECT 1")
149136

150-
# Verify that cursor.execute() set up the result set correctly
151-
self.assertIsInstance(cursor.active_result_set, ThriftResultSet)
152-
self.assertEqual(
153-
cursor.active_result_set.has_been_closed_server_side, closed
154-
)
137+
# Verify that cursor.execute() set up the result set correctly
138+
self.assertIsInstance(cursor.active_result_set, ThriftResultSet)
139+
self.assertEqual(
140+
cursor.active_result_set.has_been_closed_server_side, closed
141+
)
142+
143+
# Close the connection - this should trigger the real close chain:
144+
# connection.close() -> cursor.close() -> result_set.close()
145+
connection.close()
146+
147+
# Verify the REAL close logic worked through the chain:
148+
# 1. has_been_closed_server_side should always be True after close()
149+
self.assertTrue(real_result_set.has_been_closed_server_side)
155150

156-
# Close the connection - this should trigger the real close chain:
157-
# connection.close() -> cursor.close() -> result_set.close()
158-
connection.close()
159-
160-
# Verify the REAL close logic worked through the chain:
161-
# 1. has_been_closed_server_side should always be True after close()
162-
self.assertTrue(real_result_set.has_been_closed_server_side)
163-
164-
# 2. op_state should always be CLOSED after close()
165-
self.assertEqual(real_result_set.op_state, CommandState.CLOSED)
166-
167-
# 3. Backend close_command should be called appropriately
168-
if not closed:
169-
# Should have called backend.close_command during the close chain
170-
mock_backend.close_command.assert_called_once_with(
171-
mock_execute_response.command_id
172-
)
173-
else:
174-
# Should NOT have called backend.close_command (already closed)
175-
mock_backend.close_command.assert_not_called()
151+
# 2. op_state should always be CLOSED after close()
152+
self.assertEqual(real_result_set.op_state, CommandState.CLOSED)
153+
154+
# 3. Backend close_command should be called appropriately
155+
if not closed:
156+
# Should have called backend.close_command during the close chain
157+
mock_backend.close_command.assert_called_once_with(
158+
mock_execute_response.command_id
159+
)
160+
else:
161+
# Should NOT have called backend.close_command (already closed)
162+
mock_backend.close_command.assert_not_called()
176163

177164
@patch("%s.session.ThriftDatabricksClient" % PACKAGE_NAME)
178165
def test_cant_open_cursor_on_closed_connection(self, mock_client_class):

0 commit comments

Comments
 (0)