From 361bab4ce15ac00828236337a9b022711d4ece21 Mon Sep 17 00:00:00 2001
From: Eduard Tudenhoefner <etudenhoefner@gmail.com>
Date: Fri, 12 Apr 2024 17:45:16 +0200
Subject: [PATCH] Spark: Improvements around test initialization

---
 .../apache/iceberg/spark/extensions/ExtensionsTestBase.java   | 3 +++
 .../test/java/org/apache/iceberg/spark/CatalogTestBase.java   | 4 ----
 .../java/org/apache/iceberg/spark/TestBaseWithCatalog.java    | 2 +-
 .../java/org/apache/iceberg/spark/TestSparkExecutorCache.java | 4 ++--
 4 files changed, 6 insertions(+), 7 deletions(-)

diff --git a/spark/v3.5/spark-extensions/src/test/java/org/apache/iceberg/spark/extensions/ExtensionsTestBase.java b/spark/v3.5/spark-extensions/src/test/java/org/apache/iceberg/spark/extensions/ExtensionsTestBase.java
index 8e167b7f7320..005d5e29d5ff 100644
--- a/spark/v3.5/spark-extensions/src/test/java/org/apache/iceberg/spark/extensions/ExtensionsTestBase.java
+++ b/spark/v3.5/spark-extensions/src/test/java/org/apache/iceberg/spark/extensions/ExtensionsTestBase.java
@@ -28,6 +28,7 @@
 import org.apache.iceberg.relocated.com.google.common.collect.ImmutableMap;
 import org.apache.iceberg.spark.CatalogTestBase;
 import org.apache.iceberg.spark.TestBase;
+import org.apache.spark.api.java.JavaSparkContext;
 import org.apache.spark.sql.SparkSession;
 import org.apache.spark.sql.internal.SQLConf;
 import org.junit.jupiter.api.BeforeAll;
@@ -57,6 +58,8 @@ public static void startMetastoreAndSpark() {
             .enableHiveSupport()
             .getOrCreate();
 
+    TestBase.sparkContext = JavaSparkContext.fromSparkContext(spark.sparkContext());
+
     TestBase.catalog =
         (HiveCatalog)
             CatalogUtil.loadCatalog(
diff --git a/spark/v3.5/spark/src/test/java/org/apache/iceberg/spark/CatalogTestBase.java b/spark/v3.5/spark/src/test/java/org/apache/iceberg/spark/CatalogTestBase.java
index 6203bf89bfa1..ba864bf89e33 100644
--- a/spark/v3.5/spark/src/test/java/org/apache/iceberg/spark/CatalogTestBase.java
+++ b/spark/v3.5/spark/src/test/java/org/apache/iceberg/spark/CatalogTestBase.java
@@ -18,11 +18,9 @@
  */
 package org.apache.iceberg.spark;
 
-import java.nio.file.Path;
 import org.apache.iceberg.ParameterizedTestExtension;
 import org.apache.iceberg.Parameters;
 import org.junit.jupiter.api.extension.ExtendWith;
-import org.junit.jupiter.api.io.TempDir;
 
 @ExtendWith(ParameterizedTestExtension.class)
 public abstract class CatalogTestBase extends TestBaseWithCatalog {
@@ -48,6 +46,4 @@ protected static Object[][] parameters() {
       }
     };
   }
-
-  @TempDir protected Path temp;
 }
diff --git a/spark/v3.5/spark/src/test/java/org/apache/iceberg/spark/TestBaseWithCatalog.java b/spark/v3.5/spark/src/test/java/org/apache/iceberg/spark/TestBaseWithCatalog.java
index c3c958abf0fd..c869c4a30a19 100644
--- a/spark/v3.5/spark/src/test/java/org/apache/iceberg/spark/TestBaseWithCatalog.java
+++ b/spark/v3.5/spark/src/test/java/org/apache/iceberg/spark/TestBaseWithCatalog.java
@@ -73,7 +73,7 @@ public static void dropWarehouse() throws IOException {
     }
   }
 
-  @TempDir protected File temp;
+  @TempDir protected java.nio.file.Path temp;
 
   @Parameter(index = 0)
   protected String catalogName;
diff --git a/spark/v3.5/spark/src/test/java/org/apache/iceberg/spark/TestSparkExecutorCache.java b/spark/v3.5/spark/src/test/java/org/apache/iceberg/spark/TestSparkExecutorCache.java
index 189c125b4e3b..c5463b4249d6 100644
--- a/spark/v3.5/spark/src/test/java/org/apache/iceberg/spark/TestSparkExecutorCache.java
+++ b/spark/v3.5/spark/src/test/java/org/apache/iceberg/spark/TestSparkExecutorCache.java
@@ -410,13 +410,13 @@ private DeleteFile writeEqDeletes(Table table, String col, Object... values) thr
       deletes.add(delete.copy(col, value));
     }
 
-    OutputFile out = Files.localOutput(new File(temp, "eq-deletes-" + UUID.randomUUID()));
+    OutputFile out = Files.localOutput(new File(temp.toFile(), "eq-deletes-" + UUID.randomUUID()));
     return FileHelpers.writeDeleteFile(table, out, null, deletes, deleteSchema);
   }
 
   private Pair<DeleteFile, CharSequenceSet> writePosDeletes(
       Table table, List<Pair<CharSequence, Long>> deletes) throws IOException {
-    OutputFile out = Files.localOutput(new File(temp, "pos-deletes-" + UUID.randomUUID()));
+    OutputFile out = Files.localOutput(new File(temp.toFile(), "pos-deletes-" + UUID.randomUUID()));
     return FileHelpers.writeDeleteFile(table, out, null, deletes);
   }