Clean up tests a bit

delta-io · Jan 26, 2025 · 509b361 · 509b361
1 parent 3a97321
commit 509b361
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 27 deletions.
diff --git a/kernel/kernel-api/src/test/scala/io/delta/kernel/internal/SnapshotManagerSuite.scala b/kernel/kernel-api/src/test/scala/io/delta/kernel/internal/SnapshotManagerSuite.scala
@@ -21,7 +21,6 @@ import scala.collection.JavaConverters._
 import scala.reflect.ClassTag
 
 import io.delta.kernel.data.{ColumnVector, ColumnarBatch}
-import io.delta.kernel.engine.JsonHandler
 import io.delta.kernel.exceptions.{InvalidTableException, TableNotFoundException}
 import io.delta.kernel.expressions.Predicate
 import io.delta.kernel.internal.checkpoints.{CheckpointInstance, CheckpointMetaData, SidecarFile}
@@ -311,21 +310,9 @@ class SnapshotManagerSuite extends AnyFunSuite with MockFileSystemClientUtils {
       lastCheckpointVersion: Optional[java.lang.Long] = Optional.empty(),
       versionToLoad: Optional[java.lang.Long] = Optional.empty(),
       expectedErrorMessageContains: String = "")(implicit classTag: ClassTag[T]): Unit = {
-    val jsonHandler = if (lastCheckpointVersion.isPresent) {
-      new MockReadLastCheckpointFileJsonHandler(
-        s"$logPath/_last_checkpoint",
-        lastCheckpointVersion.get()
-      )
-    } else {
-      null
-    }
-
     val e = intercept[T] {
       snapshotManager.getLogSegmentForVersion(
-        mockEngine(
-          jsonHandler = jsonHandler,
-          fileSystemClient = new MockListFromFileSystemClient(listFromProvider(files))
-        ),
+        createMockFSAndJsonEngineForLastCheckpoint(files, lastCheckpointVersion),
         versionToLoad
       )
     }
@@ -791,14 +778,6 @@ class SnapshotManagerSuite extends AnyFunSuite with MockFileSystemClientUtils {
       (10, Seq.empty, (0L to 10L), Optional.empty())
     )
     cases.foreach { case (corruptedVersion, validVersions, deltaVersions, lastCheckpointVersion) =>
-      val jsonHandler = if (lastCheckpointVersion.isPresent) {
-        new MockReadLastCheckpointFileJsonHandler(
-          s"$logPath/_last_checkpoint",
-          lastCheckpointVersion.get()
-        )
-      } else {
-        null
-      }
       val corruptedCheckpoint = FileNames.checkpointFileWithParts(logPath, corruptedVersion, 5)
         .asScala
         .map(p => FileStatus.of(p.toString, 10, 10))
@@ -807,10 +786,7 @@ class SnapshotManagerSuite extends AnyFunSuite with MockFileSystemClientUtils {
       val deltas = deltaFileStatuses(deltaVersions)
       val allFiles = deltas ++ corruptedCheckpoint ++ checkpoints
       val logSegment = snapshotManager.getLogSegmentForVersion(
-        mockEngine(
-          jsonHandler = jsonHandler,
-          fileSystemClient = new MockListFromFileSystemClient(listFromProvider(allFiles))
-        ),
+        createMockFSAndJsonEngineForLastCheckpoint(allFiles, lastCheckpointVersion),
         Optional.empty()
       )
       val checkpointVersion = validVersions.sorted.lastOption

diff --git a/kernel/kernel-api/src/test/scala/io/delta/kernel/test/MockFileSystemClientUtils.scala b/kernel/kernel-api/src/test/scala/io/delta/kernel/test/MockFileSystemClientUtils.scala
@@ -15,7 +15,7 @@
  */
 package io.delta.kernel.test
 
-import java.util.UUID
+import java.util.{Optional, UUID}
 
 import io.delta.kernel.engine._
 import io.delta.kernel.internal.fs.Path
@@ -25,6 +25,8 @@ import io.delta.kernel.utils.{CloseableIterator, FileStatus}
 
 import scala.collection.JavaConverters._
 
+import io.delta.kernel.internal.MockReadLastCheckpointFileJsonHandler
+
 /**
  * This is an extension to [[BaseMockFileSystemClient]] containing specific mock implementations
  * [[FileSystemClient]] which are shared across multiple test suite.
@@ -107,6 +109,21 @@ trait MockFileSystemClientUtils extends MockEngineUtils {
       jsonHandler = jsonHandler)
   }
 
+  def createMockFSAndJsonEngineForLastCheckpoint(
+      contents: Seq[FileStatus], lastCheckpointVersion: Optional[java.lang.Long]): Engine = {
+    mockEngine(
+      fileSystemClient = new MockListFromFileSystemClient(listFromProvider(contents)),
+      jsonHandler = if (lastCheckpointVersion.isPresent) {
+        new MockReadLastCheckpointFileJsonHandler(
+          s"$logPath/_last_checkpoint",
+          lastCheckpointVersion.get()
+        )
+      } else {
+        null
+      }
+    )
+  }
+
   /**
    * Create a mock [[Engine]] to mock the [[FileSystemClient.listFrom]] calls using
    * the given contents. The contents are filtered depending upon the list from path prefix.