update unit tests to address ThriftBackend through session instead of through Connection

Signed-off-by: varun-edachali-dbx <[email protected]>

Author: varun-edachali-dbx

tests/unit/test_client.py

@@ -80,7 +80,7 @@ class ClientTestSuite(unittest.TestCase):
         "access_token": "tok",
     }
 
-    @patch("%s.client.ThriftBackend" % PACKAGE_NAME)
+    @patch("%s.session.ThriftBackend" % PACKAGE_NAME)
     def test_close_uses_the_correct_session_id(self, mock_client_class):
         instance = mock_client_class.return_value
 
@@ -95,7 +95,7 @@ def test_close_uses_the_correct_session_id(self, mock_client_class):
         close_session_id = instance.close_session.call_args[0][0].sessionId
         self.assertEqual(close_session_id, b"\x22")
 
-    @patch("%s.client.ThriftBackend" % PACKAGE_NAME)
+    @patch("%s.session.ThriftBackend" % PACKAGE_NAME)
     def test_auth_args(self, mock_client_class):
         # Test that the following auth args work:
         # token = foo,
@@ -122,15 +122,15 @@ def test_auth_args(self, mock_client_class):
             self.assertEqual(args["http_path"], http_path)
             connection.close()
 
-    @patch("%s.client.ThriftBackend" % PACKAGE_NAME)
+    @patch("%s.session.ThriftBackend" % PACKAGE_NAME)
     def test_http_header_passthrough(self, mock_client_class):
         http_headers = [("foo", "bar")]
         databricks.sql.connect(**self.DUMMY_CONNECTION_ARGS, http_headers=http_headers)
 
         call_args = mock_client_class.call_args[0][3]
         self.assertIn(("foo", "bar"), call_args)
 
-    @patch("%s.client.ThriftBackend" % PACKAGE_NAME)
+    @patch("%s.session.ThriftBackend" % PACKAGE_NAME)
     def test_tls_arg_passthrough(self, mock_client_class):
         databricks.sql.connect(
             **self.DUMMY_CONNECTION_ARGS,
@@ -146,7 +146,7 @@ def test_tls_arg_passthrough(self, mock_client_class):
         self.assertEqual(kwargs["_tls_client_cert_key_file"], "trusted client cert")
         self.assertEqual(kwargs["_tls_client_cert_key_password"], "key password")
 
-    @patch("%s.client.ThriftBackend" % PACKAGE_NAME)
+    @patch("%s.session.ThriftBackend" % PACKAGE_NAME)
     def test_useragent_header(self, mock_client_class):
         databricks.sql.connect(**self.DUMMY_CONNECTION_ARGS)
 
@@ -167,7 +167,7 @@ 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.session.ThriftBackend" % PACKAGE_NAME, ThriftBackendMockFactory.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
@@ -184,7 +184,7 @@ def test_closing_connection_closes_commands(self, mock_result_set_class):
                 )
                 mock_result_set_class.return_value.close.assert_called_once_with()
 
-    @patch("%s.client.ThriftBackend" % PACKAGE_NAME)
+    @patch("%s.session.ThriftBackend" % PACKAGE_NAME)
     def test_cant_open_cursor_on_closed_connection(self, mock_client_class):
         connection = databricks.sql.connect(**self.DUMMY_CONNECTION_ARGS)
         self.assertTrue(connection.open)
@@ -194,7 +194,7 @@ def test_cant_open_cursor_on_closed_connection(self, mock_client_class):
             connection.cursor()
         self.assertIn("closed", str(cm.exception))
 
-    @patch("%s.client.ThriftBackend" % PACKAGE_NAME)
+    @patch("%s.session.ThriftBackend" % PACKAGE_NAME)
     @patch("%s.client.Cursor" % PACKAGE_NAME)
     def test_arraysize_buffer_size_passthrough(
         self, mock_cursor_class, mock_client_class
@@ -214,7 +214,10 @@ def test_closing_result_set_with_closed_connection_soft_closes_commands(self):
             thrift_backend=mock_backend,
             execute_response=Mock(),
         )
-        mock_connection.open = False
+        # Setup session mock on the mock_connection
+        mock_session = Mock()
+        mock_session.open = False
+        type(mock_connection).session = PropertyMock(return_value=mock_session)
 
         result_set.close()
 
@@ -226,7 +229,11 @@ def test_closing_result_set_hard_closes_commands(self):
         mock_results_response.has_been_closed_server_side = False
         mock_connection = Mock()
         mock_thrift_backend = Mock()
-        mock_connection.open = True
+        # Setup session mock on the mock_connection
+        mock_session = Mock()
+        mock_session.open = True
+        type(mock_connection).session = PropertyMock(return_value=mock_session)
+        
         result_set = client.ResultSet(
             mock_connection, mock_results_response, mock_thrift_backend
         )
@@ -283,7 +290,7 @@ def test_context_manager_closes_cursor(self):
             cursor.close = mock_close
         mock_close.assert_called_once_with()
 
-    @patch("%s.client.ThriftBackend" % PACKAGE_NAME)
+    @patch("%s.session.ThriftBackend" % PACKAGE_NAME)
     def test_context_manager_closes_connection(self, mock_client_class):
         instance = mock_client_class.return_value
 
@@ -396,7 +403,7 @@ def test_cancel_command_will_issue_warning_for_cancel_with_no_executing_command(
         self.assertTrue(logger_instance.warning.called)
         self.assertFalse(mock_thrift_backend.cancel_command.called)
 
-    @patch("%s.client.ThriftBackend" % PACKAGE_NAME)
+    @patch("%s.session.ThriftBackend" % PACKAGE_NAME)
     def test_max_number_of_retries_passthrough(self, mock_client_class):
         databricks.sql.connect(
             _retry_stop_after_attempts_count=54, **self.DUMMY_CONNECTION_ARGS
@@ -406,7 +413,7 @@ def test_max_number_of_retries_passthrough(self, mock_client_class):
             mock_client_class.call_args[1]["_retry_stop_after_attempts_count"], 54
         )
 
-    @patch("%s.client.ThriftBackend" % PACKAGE_NAME)
+    @patch("%s.session.ThriftBackend" % PACKAGE_NAME)
     def test_socket_timeout_passthrough(self, mock_client_class):
         databricks.sql.connect(_socket_timeout=234, **self.DUMMY_CONNECTION_ARGS)
         self.assertEqual(mock_client_class.call_args[1]["_socket_timeout"], 234)
@@ -419,7 +426,7 @@ def test_version_is_canonical(self):
         )
         self.assertIsNotNone(re.match(canonical_version_re, version))
 
-    @patch("%s.client.ThriftBackend" % PACKAGE_NAME)
+    @patch("%s.session.ThriftBackend" % PACKAGE_NAME)
     def test_configuration_passthrough(self, mock_client_class):
         mock_session_config = Mock()
         databricks.sql.connect(
@@ -431,7 +438,7 @@ def test_configuration_passthrough(self, mock_client_class):
             mock_session_config,
         )
 
-    @patch("%s.client.ThriftBackend" % PACKAGE_NAME)
+    @patch("%s.session.ThriftBackend" % PACKAGE_NAME)
     def test_initial_namespace_passthrough(self, mock_client_class):
         mock_cat = Mock()
         mock_schem = Mock()
@@ -505,7 +512,7 @@ def test_executemany_parameter_passhthrough_and_uses_last_result_set(
             "last operation",
         )
 
-    @patch("%s.client.ThriftBackend" % PACKAGE_NAME)
+    @patch("%s.session.ThriftBackend" % PACKAGE_NAME)
     def test_commit_a_noop(self, mock_thrift_backend_class):
         c = databricks.sql.connect(**self.DUMMY_CONNECTION_ARGS)
         c.commit()
@@ -518,7 +525,7 @@ def test_setoutputsizes_a_noop(self):
         cursor = client.Cursor(Mock(), Mock())
         cursor.setoutputsize(1)
 
-    @patch("%s.client.ThriftBackend" % PACKAGE_NAME)
+    @patch("%s.session.ThriftBackend" % PACKAGE_NAME)
     def test_rollback_not_supported(self, mock_thrift_backend_class):
         c = databricks.sql.connect(**self.DUMMY_CONNECTION_ARGS)
         with self.assertRaises(NotSupportedError):
@@ -603,7 +610,7 @@ def test_column_name_api(self):
                 },
             )
 
-    @patch("%s.client.ThriftBackend" % PACKAGE_NAME)
+    @patch("%s.session.ThriftBackend" % PACKAGE_NAME)
     def test_finalizer_closes_abandoned_connection(self, mock_client_class):
         instance = mock_client_class.return_value
 
@@ -620,7 +627,7 @@ def test_finalizer_closes_abandoned_connection(self, mock_client_class):
         close_session_id = instance.close_session.call_args[0][0].sessionId
         self.assertEqual(close_session_id, b"\x22")
 
-    @patch("%s.client.ThriftBackend" % PACKAGE_NAME)
+    @patch("%s.session.ThriftBackend" % PACKAGE_NAME)
     def test_cursor_keeps_connection_alive(self, mock_client_class):
         instance = mock_client_class.return_value
 
@@ -639,7 +646,7 @@ def test_cursor_keeps_connection_alive(self, mock_client_class):
 
     @patch("%s.utils.ExecuteResponse" % PACKAGE_NAME, autospec=True)
     @patch("%s.client.Cursor._handle_staging_operation" % PACKAGE_NAME)
-    @patch("%s.client.ThriftBackend" % PACKAGE_NAME)
+    @patch("%s.session.ThriftBackend" % PACKAGE_NAME)
     def test_staging_operation_response_is_handled(
         self, mock_client_class, mock_handle_staging_operation, mock_execute_response
     ):
@@ -658,7 +665,7 @@ def test_staging_operation_response_is_handled(
 
         mock_handle_staging_operation.call_count == 1
 
-    @patch("%s.client.ThriftBackend" % PACKAGE_NAME, ThriftBackendMockFactory.new())
+    @patch("%s.session.ThriftBackend" % PACKAGE_NAME, ThriftBackendMockFactory.new())
     def test_access_current_query_id(self):
         operation_id = "EE6A8778-21FC-438B-92D8-96AC51EE3821"