[SPARK-58407][SQL][TESTS][CONNECT] Fix connect.SessionQueryTest beforeAll, afterAll

sql/connect/server/src/test/scala/org/apache/spark/sql/connect/SessionQueryTestBeforeAfterHooksSuite.scala

@@ -0,0 +1,24 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.connect
+
+import org.apache.spark.sql
+
+class SessionQueryTestBeforeAfterHooksSuite
+    extends sql.SessionQueryTestBeforeAfterHooksSuite
+    with SessionQueryTest

sql/connect/server/src/test/scala/org/apache/spark/sql/connect/SparkSessionBinder.scala

@@ -38,7 +38,8 @@ trait SparkSessionBinder extends sql.SparkSessionBinder { self: SparkFunSuite =>
   protected override def spark: SparkSession = _connectSpark
 
   override protected def beforeAll(): Unit = {
-    super.beforeAll()
+    initializeSession()
+
     // Other suites using mocks leave a mess in the global executionManager,
     // shut it down so that it's cleared before starting server.
     SparkConnectService.executionManager.shutdown()
@@ -60,14 +61,18 @@ trait SparkSessionBinder extends sql.SparkSessionBinder { self: SparkFunSuite =>
       .builder()
       .client(client)
       .create()
+    super.beforeAll()
   }
 
   override def afterAll(): Unit = {
-    if (_connectSpark != null) {
-      _connectSpark.close()
-      _connectSpark = null
+    try {
+      super.afterAll()
+    } finally {
+      if (_connectSpark != null) {
+        _connectSpark.close()
+        _connectSpark = null
+      }
+      SparkConnectService.stop()
     }
-    SparkConnectService.stop()
-    super.afterAll()
   }
 }

sql/core/src/test/scala/org/apache/spark/sql/SessionQueryTestBeforeAfterHooksSuite.scala

@@ -0,0 +1,47 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql
+
+class SessionQueryTestBeforeAfterHooksSuite extends SessionQueryTest {
+
+  override protected def beforeAll(): Unit = {
+    super.beforeAll()
+    checkAnswer(spark.sql("SELECT 1"), Seq(Row(1)))
+  }
+
+  override protected def beforeEach(): Unit = {
+    checkAnswer(spark.sql("SELECT 1"), Seq(Row(1)))
+    super.beforeEach()
+    checkAnswer(spark.sql("SELECT 1"), Seq(Row(1)))
+  }
+
+  override protected def afterEach(): Unit = {
+    checkAnswer(spark.sql("SELECT 1"), Seq(Row(1)))
+    super.afterEach()
+    checkAnswer(spark.sql("SELECT 1"), Seq(Row(1)))
+  }
+
+  override protected def afterAll(): Unit = {
+    checkAnswer(spark.sql("SELECT 1"), Seq(Row(1)))
+    super.afterAll()
+  }
+
+  test("assert spark is available in BeforeAndAfter hooks") {
+    checkAnswer(spark.sql("SELECT 1"), Seq(Row(1)))
+  }
+}