From 47f4547a0b3f8cb5af9598937979627b2c86561f Mon Sep 17 00:00:00 2001 From: varun-edachali-dbx Date: Wed, 4 Jun 2025 01:46:48 +0000 Subject: [PATCH 1/7] [init] working code for corrected test Signed-off-by: varun-edachali-dbx --- tests/unit/test_client.py | 70 ++++++++++++++++++++++++++++++++++----- 1 file changed, 62 insertions(+), 8 deletions(-) diff --git a/tests/unit/test_client.py b/tests/unit/test_client.py index 5271baa70..553d5b21d 100644 --- a/tests/unit/test_client.py +++ b/tests/unit/test_client.py @@ -13,6 +13,7 @@ TExecuteStatementResp, TOperationHandle, THandleIdentifier, + TOperationState, TOperationType, ) from databricks.sql.thrift_backend import ThriftBackend @@ -23,6 +24,7 @@ from databricks.sql.exc import RequestError, CursorAlreadyClosedError from databricks.sql.types import Row +from databricks.sql.utils import ExecuteResponse from tests.unit.test_fetches import FetchTests from tests.unit.test_thrift_backend import ThriftBackendTestSuite from tests.unit.test_arrow_queue import ArrowQueueSuite @@ -168,22 +170,74 @@ def test_useragent_header(self, mock_client_class): http_headers = mock_client_class.call_args[0][3] self.assertIn(user_agent_header_with_entry, http_headers) - @patch("%s.client.ThriftBackend" % PACKAGE_NAME, ThriftBackendMockFactory.new()) - @patch("%s.client.ResultSet" % PACKAGE_NAME) - def test_closing_connection_closes_commands(self, mock_result_set_class): + @patch("%s.client.ThriftBackend" % PACKAGE_NAME) + def test_closing_connection_closes_commands(self, mock_thrift_client_class): # Test once with has_been_closed_server side, once without for closed in (True, False): with self.subTest(closed=closed): - mock_result_set_class.return_value = Mock() + initial_state = ( + TOperationState.FINISHED_STATE + if not closed + else TOperationState.CLOSED_STATE + ) + + # Mock the execute response with controlled state + mock_execute_response = Mock(spec=ExecuteResponse) + mock_execute_response.status = initial_state + mock_execute_response.has_been_closed_server_side = closed + mock_execute_response.is_staging_operation = False + + # Mock the backend that will be used by the real ThriftResultSet + mock_backend = Mock(spec=ThriftBackend) + + # Configure the decorator's mock to return our specific mock_backend + mock_thrift_client_class.return_value = mock_backend + + # Create connection and cursor connection = databricks.sql.connect(**self.DUMMY_CONNECTION_ARGS) cursor = connection.cursor() - cursor.execute("SELECT 1;") + + # Verify initial state + self.assertEqual( + mock_execute_response.has_been_closed_server_side, closed + ) + self.assertEqual(mock_execute_response.status, initial_state) + + # Mock execute_command to return our execute response + cursor.thrift_backend.execute_command = Mock( + return_value=mock_execute_response + ) + + # Execute a command - this should set cursor.active_result_set to our real result set + cursor.execute("SELECT 1") + + # Verify that cursor.execute() set up the result set correctly + active_result_set = cursor.active_result_set + self.assertEqual(active_result_set.has_been_closed_server_side, closed) + + # Close the connection - this should trigger the real close chain: + # connection.close() -> cursor.close() -> result_set.close() connection.close() - self.assertTrue( - mock_result_set_class.return_value.has_been_closed_server_side + # Verify the REAL close logic worked through the chain: + # 1. has_been_closed_server_side should always be True after close() + self.assertTrue(active_result_set.has_been_closed_server_side) + + # 2. op_state should always be CLOSED after close() + self.assertEqual( + active_result_set.op_state, + connection.thrift_backend.CLOSED_OP_STATE, ) - mock_result_set_class.return_value.close.assert_called_once_with() + + # 3. Backend close_command should be called appropriately + if not closed: + # Should have called backend.close_command during the close chain + mock_backend.close_command.assert_called_once_with( + mock_execute_response.command_handle + ) + else: + # Should NOT have called backend.close_command (already closed) + mock_backend.close_command.assert_not_called() @patch("%s.client.ThriftBackend" % PACKAGE_NAME) def test_cant_open_cursor_on_closed_connection(self, mock_client_class): From 2f666fb15075675c55c2182b9bf47f4c24b521af Mon Sep 17 00:00:00 2001 From: varun-edachali-dbx Date: Wed, 4 Jun 2025 04:48:23 +0000 Subject: [PATCH 2/7] excess comments Signed-off-by: varun-edachali-dbx --- tests/unit/test_client.py | 11 +++-------- 1 file changed, 3 insertions(+), 8 deletions(-) diff --git a/tests/unit/test_client.py b/tests/unit/test_client.py index 553d5b21d..221f85c8f 100644 --- a/tests/unit/test_client.py +++ b/tests/unit/test_client.py @@ -187,10 +187,8 @@ def test_closing_connection_closes_commands(self, mock_thrift_client_class): mock_execute_response.has_been_closed_server_side = closed mock_execute_response.is_staging_operation = False - # Mock the backend that will be used by the real ThriftResultSet + # Mock the backend that will be used mock_backend = Mock(spec=ThriftBackend) - - # Configure the decorator's mock to return our specific mock_backend mock_thrift_client_class.return_value = mock_backend # Create connection and cursor @@ -207,19 +205,16 @@ def test_closing_connection_closes_commands(self, mock_thrift_client_class): cursor.thrift_backend.execute_command = Mock( return_value=mock_execute_response ) - - # Execute a command - this should set cursor.active_result_set to our real result set cursor.execute("SELECT 1") # Verify that cursor.execute() set up the result set correctly active_result_set = cursor.active_result_set self.assertEqual(active_result_set.has_been_closed_server_side, closed) - # Close the connection - this should trigger the real close chain: - # connection.close() -> cursor.close() -> result_set.close() + # Close the connection connection.close() - # Verify the REAL close logic worked through the chain: + # Verify the close logic worked: # 1. has_been_closed_server_side should always be True after close() self.assertTrue(active_result_set.has_been_closed_server_side) From f96936ac45b8ab413f701eb5cee3af9e15b547eb Mon Sep 17 00:00:00 2001 From: varun-edachali-dbx Date: Mon, 9 Jun 2025 12:30:58 +0000 Subject: [PATCH 3/7] use pytest parametrize instead of unittest subtests, switch to pytest instead of unittest primitives Signed-off-by: varun-edachali-dbx --- tests/unit/test_client.py | 135 ++++++++++++++++++++------------------ 1 file changed, 71 insertions(+), 64 deletions(-) diff --git a/tests/unit/test_client.py b/tests/unit/test_client.py index 221f85c8f..8f4bf15f6 100644 --- a/tests/unit/test_client.py +++ b/tests/unit/test_client.py @@ -2,6 +2,7 @@ import re import sys import unittest +import pytest from unittest.mock import patch, MagicMock, Mock, PropertyMock import itertools from decimal import Decimal @@ -30,6 +31,76 @@ from tests.unit.test_arrow_queue import ArrowQueueSuite +@pytest.mark.parametrize("closed", [True, False]) +@patch("databricks.sql.client.ThriftBackend") +def test_closing_connection_closes_commands(mock_thrift_client_class, closed): + """Test that closing a connection properly closes commands. + + This test verifies that when a connection is closed: + 1. All result sets are marked as closed server-side + 2. The operation state is set to CLOSED + 3. Backend.close_command is called only for commands that weren't already closed + + Args: + mock_thrift_client_class: Mock for ThriftBackend class + closed: Parameter indicating if the command is already closed + """ + # Set initial state based on whether the command is already closed + initial_state = ( + TOperationState.FINISHED_STATE if not closed else TOperationState.CLOSED_STATE + ) + + # Mock the execute response with controlled state + mock_execute_response = Mock(spec=ExecuteResponse) + mock_execute_response.status = initial_state + mock_execute_response.has_been_closed_server_side = closed + mock_execute_response.is_staging_operation = False + + # Mock the backend that will be used + mock_backend = Mock(spec=ThriftBackend) + mock_thrift_client_class.return_value = mock_backend + + # Create connection and cursor + connection = databricks.sql.connect( + server_hostname="foo", + http_path="dummy_path", + access_token="tok", + ) + cursor = connection.cursor() + + # Verify initial state + assert mock_execute_response.has_been_closed_server_side == closed + assert mock_execute_response.status == initial_state + + # Mock execute_command to return our execute response + cursor.thrift_backend.execute_command = Mock(return_value=mock_execute_response) + cursor.execute("SELECT 1") + + # Verify that cursor.execute() set up the result set correctly + active_result_set = cursor.active_result_set + assert active_result_set.has_been_closed_server_side == closed + + # Close the connection + connection.close() + + # Verify the close logic worked: + # 1. has_been_closed_server_side should always be True after close() + assert active_result_set.has_been_closed_server_side is True + + # 2. op_state should always be CLOSED after close() + assert active_result_set.op_state == connection.thrift_backend.CLOSED_OP_STATE + + # 3. Backend close_command should be called appropriately + if not closed: + # Should have called backend.close_command during the close chain + mock_backend.close_command.assert_called_once_with( + mock_execute_response.command_handle + ) + else: + # Should NOT have called backend.close_command (already closed) + mock_backend.close_command.assert_not_called() + + class ThriftBackendMockFactory: @classmethod def new(cls): @@ -170,70 +241,6 @@ def test_useragent_header(self, mock_client_class): http_headers = mock_client_class.call_args[0][3] self.assertIn(user_agent_header_with_entry, http_headers) - @patch("%s.client.ThriftBackend" % PACKAGE_NAME) - def test_closing_connection_closes_commands(self, mock_thrift_client_class): - # Test once with has_been_closed_server side, once without - for closed in (True, False): - with self.subTest(closed=closed): - initial_state = ( - TOperationState.FINISHED_STATE - if not closed - else TOperationState.CLOSED_STATE - ) - - # Mock the execute response with controlled state - mock_execute_response = Mock(spec=ExecuteResponse) - mock_execute_response.status = initial_state - mock_execute_response.has_been_closed_server_side = closed - mock_execute_response.is_staging_operation = False - - # Mock the backend that will be used - mock_backend = Mock(spec=ThriftBackend) - mock_thrift_client_class.return_value = mock_backend - - # Create connection and cursor - connection = databricks.sql.connect(**self.DUMMY_CONNECTION_ARGS) - cursor = connection.cursor() - - # Verify initial state - self.assertEqual( - mock_execute_response.has_been_closed_server_side, closed - ) - self.assertEqual(mock_execute_response.status, initial_state) - - # Mock execute_command to return our execute response - cursor.thrift_backend.execute_command = Mock( - return_value=mock_execute_response - ) - cursor.execute("SELECT 1") - - # Verify that cursor.execute() set up the result set correctly - active_result_set = cursor.active_result_set - self.assertEqual(active_result_set.has_been_closed_server_side, closed) - - # Close the connection - connection.close() - - # Verify the close logic worked: - # 1. has_been_closed_server_side should always be True after close() - self.assertTrue(active_result_set.has_been_closed_server_side) - - # 2. op_state should always be CLOSED after close() - self.assertEqual( - active_result_set.op_state, - connection.thrift_backend.CLOSED_OP_STATE, - ) - - # 3. Backend close_command should be called appropriately - if not closed: - # Should have called backend.close_command during the close chain - mock_backend.close_command.assert_called_once_with( - mock_execute_response.command_handle - ) - else: - # Should NOT have called backend.close_command (already closed) - mock_backend.close_command.assert_not_called() - @patch("%s.client.ThriftBackend" % PACKAGE_NAME) def test_cant_open_cursor_on_closed_connection(self, mock_client_class): connection = databricks.sql.connect(**self.DUMMY_CONNECTION_ARGS) From 5521f21283845d65cfc27eb8e372ba3184c29542 Mon Sep 17 00:00:00 2001 From: varun-edachali-dbx Date: Mon, 9 Jun 2025 12:37:08 +0000 Subject: [PATCH 4/7] initialisations, function calls, then assertions Signed-off-by: varun-edachali-dbx --- tests/unit/test_client.py | 13 +++++-------- 1 file changed, 5 insertions(+), 8 deletions(-) diff --git a/tests/unit/test_client.py b/tests/unit/test_client.py index 8f4bf15f6..c08b73fdb 100644 --- a/tests/unit/test_client.py +++ b/tests/unit/test_client.py @@ -68,18 +68,15 @@ def test_closing_connection_closes_commands(mock_thrift_client_class, closed): ) cursor = connection.cursor() - # Verify initial state - assert mock_execute_response.has_been_closed_server_side == closed - assert mock_execute_response.status == initial_state - # Mock execute_command to return our execute response cursor.thrift_backend.execute_command = Mock(return_value=mock_execute_response) - cursor.execute("SELECT 1") - # Verify that cursor.execute() set up the result set correctly + # Execute a command + cursor.execute("SELECT 1") + + # Get the active result set for later assertions active_result_set = cursor.active_result_set - assert active_result_set.has_been_closed_server_side == closed - + # Close the connection connection.close() From c22a53ff57291573da5a75bd596fd729e0d8b3d5 Mon Sep 17 00:00:00 2001 From: varun-edachali-dbx Date: Tue, 10 Jun 2025 11:34:42 +0530 Subject: [PATCH 5/7] remove pytest parametrize, move back to unittest subTest to allow keeping the test inside ClientTestSuite Signed-off-by: varun-edachali-dbx --- tests/unit/test_client.py | 141 ++++++++++++++++++++------------------ 1 file changed, 74 insertions(+), 67 deletions(-) diff --git a/tests/unit/test_client.py b/tests/unit/test_client.py index c08b73fdb..bf3b170e7 100644 --- a/tests/unit/test_client.py +++ b/tests/unit/test_client.py @@ -31,73 +31,6 @@ from tests.unit.test_arrow_queue import ArrowQueueSuite -@pytest.mark.parametrize("closed", [True, False]) -@patch("databricks.sql.client.ThriftBackend") -def test_closing_connection_closes_commands(mock_thrift_client_class, closed): - """Test that closing a connection properly closes commands. - - This test verifies that when a connection is closed: - 1. All result sets are marked as closed server-side - 2. The operation state is set to CLOSED - 3. Backend.close_command is called only for commands that weren't already closed - - Args: - mock_thrift_client_class: Mock for ThriftBackend class - closed: Parameter indicating if the command is already closed - """ - # Set initial state based on whether the command is already closed - initial_state = ( - TOperationState.FINISHED_STATE if not closed else TOperationState.CLOSED_STATE - ) - - # Mock the execute response with controlled state - mock_execute_response = Mock(spec=ExecuteResponse) - mock_execute_response.status = initial_state - mock_execute_response.has_been_closed_server_side = closed - mock_execute_response.is_staging_operation = False - - # Mock the backend that will be used - mock_backend = Mock(spec=ThriftBackend) - mock_thrift_client_class.return_value = mock_backend - - # Create connection and cursor - connection = databricks.sql.connect( - server_hostname="foo", - http_path="dummy_path", - access_token="tok", - ) - cursor = connection.cursor() - - # Mock execute_command to return our execute response - cursor.thrift_backend.execute_command = Mock(return_value=mock_execute_response) - - # Execute a command - cursor.execute("SELECT 1") - - # Get the active result set for later assertions - active_result_set = cursor.active_result_set - - # Close the connection - connection.close() - - # Verify the close logic worked: - # 1. has_been_closed_server_side should always be True after close() - assert active_result_set.has_been_closed_server_side is True - - # 2. op_state should always be CLOSED after close() - assert active_result_set.op_state == connection.thrift_backend.CLOSED_OP_STATE - - # 3. Backend close_command should be called appropriately - if not closed: - # Should have called backend.close_command during the close chain - mock_backend.close_command.assert_called_once_with( - mock_execute_response.command_handle - ) - else: - # Should NOT have called backend.close_command (already closed) - mock_backend.close_command.assert_not_called() - - class ThriftBackendMockFactory: @classmethod def new(cls): @@ -238,6 +171,80 @@ def test_useragent_header(self, mock_client_class): http_headers = mock_client_class.call_args[0][3] self.assertIn(user_agent_header_with_entry, http_headers) + @patch("databricks.sql.client.ThriftBackend") + def test_closing_connection_closes_commands(self, mock_thrift_client_class): + """Test that closing a connection properly closes commands. + + This test verifies that when a connection is closed: + 1. All result sets are marked as closed server-side + 2. The operation state is set to CLOSED + 3. Backend.close_command is called only for commands that weren't already closed + + Args: + mock_thrift_client_class: Mock for ThriftBackend class + closed: Parameter indicating if the command is already closed + """ + for closed in (True, False): + with self.subTest(closed=closed): + # Set initial state based on whether the command is already closed + initial_state = ( + TOperationState.FINISHED_STATE + if not closed + else TOperationState.CLOSED_STATE + ) + + # Mock the execute response with controlled state + mock_execute_response = Mock(spec=ExecuteResponse) + mock_execute_response.status = initial_state + mock_execute_response.has_been_closed_server_side = closed + mock_execute_response.is_staging_operation = False + + # Mock the backend that will be used + mock_backend = Mock(spec=ThriftBackend) + mock_thrift_client_class.return_value = mock_backend + + # Create connection and cursor + connection = databricks.sql.connect( + server_hostname="foo", + http_path="dummy_path", + access_token="tok", + ) + cursor = connection.cursor() + + # Mock execute_command to return our execute response + cursor.thrift_backend.execute_command = Mock( + return_value=mock_execute_response + ) + + # Execute a command + cursor.execute("SELECT 1") + + # Get the active result set for later assertions + active_result_set = cursor.active_result_set + + # Close the connection + connection.close() + + # Verify the close logic worked: + # 1. has_been_closed_server_side should always be True after close() + assert active_result_set.has_been_closed_server_side is True + + # 2. op_state should always be CLOSED after close() + assert ( + active_result_set.op_state + == connection.thrift_backend.CLOSED_OP_STATE + ) + + # 3. Backend close_command should be called appropriately + if not closed: + # Should have called backend.close_command during the close chain + mock_backend.close_command.assert_called_once_with( + mock_execute_response.command_handle + ) + else: + # Should NOT have called backend.close_command (already closed) + mock_backend.close_command.assert_not_called() + @patch("%s.client.ThriftBackend" % PACKAGE_NAME) def test_cant_open_cursor_on_closed_connection(self, mock_client_class): connection = databricks.sql.connect(**self.DUMMY_CONNECTION_ARGS) From c3d8b0f66938e322f8f33d67ba2f5a8ea19746b7 Mon Sep 17 00:00:00 2001 From: varun-edachali-dbx Date: Tue, 10 Jun 2025 11:37:34 +0530 Subject: [PATCH 6/7] clean up docstring Signed-off-by: varun-edachali-dbx --- tests/unit/test_client.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/tests/unit/test_client.py b/tests/unit/test_client.py index bf3b170e7..ec91155f0 100644 --- a/tests/unit/test_client.py +++ b/tests/unit/test_client.py @@ -176,13 +176,12 @@ def test_closing_connection_closes_commands(self, mock_thrift_client_class): """Test that closing a connection properly closes commands. This test verifies that when a connection is closed: - 1. All result sets are marked as closed server-side + 1. the active result set is marked as closed server-side 2. The operation state is set to CLOSED - 3. Backend.close_command is called only for commands that weren't already closed + 3. backend.close_command is called only for commands that weren't already closed Args: mock_thrift_client_class: Mock for ThriftBackend class - closed: Parameter indicating if the command is already closed """ for closed in (True, False): with self.subTest(closed=closed): From dea05e6874d5434572c43974be1b08217792ca05 Mon Sep 17 00:00:00 2001 From: varun-edachali-dbx Date: Tue, 10 Jun 2025 13:52:36 +0530 Subject: [PATCH 7/7] remove pytest import Signed-off-by: varun-edachali-dbx --- tests/unit/test_client.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/unit/test_client.py b/tests/unit/test_client.py index ec91155f0..cf8779a21 100644 --- a/tests/unit/test_client.py +++ b/tests/unit/test_client.py @@ -2,7 +2,6 @@ import re import sys import unittest -import pytest from unittest.mock import patch, MagicMock, Mock, PropertyMock import itertools from decimal import Decimal