remove un-necessary artifacts in test_session, add back assertion

Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com>

Author: varun-edachali-dbx

tests/unit/test_session.py

@@ -128,6 +128,15 @@ def test_context_manager_closes_connection(self, mock_client_class):
         self.assertEqual(close_session_call_args.guid, b"\x22")
         self.assertEqual(close_session_call_args.secret, b"\x33")
 
+        connection = databricks.sql.connect(**self.DUMMY_CONNECTION_ARGS)
+        connection.close = Mock()
+        try:
+            with self.assertRaises(KeyboardInterrupt):
+                with connection:
+                    raise KeyboardInterrupt("Simulated interrupt")
+        finally:
+            connection.close.assert_called()
+
     @patch("%s.session.ThriftDatabricksClient" % PACKAGE_NAME)
     def test_max_number_of_retries_passthrough(self, mock_client_class):
         databricks.sql.connect(
@@ -146,33 +155,21 @@ 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
         )
 
-        # Check that open_session was called with the correct session_configuration as keyword argument
         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
         )
 
-        # Check that open_session was called with the correct catalog and schema as keyword arguments
         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)
@@ -181,7 +178,6 @@ def test_initial_namespace_passthrough(self, mock_client_class):
     def test_finalizer_closes_abandoned_connection(self, mock_client_class):
         instance = mock_client_class.return_value
 
-        # 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