From d65ae3f8abfa071241022f2c3b966a0833725208 Mon Sep 17 00:00:00 2001
From: Bolek Ziobrowski <26925920+bziobrowski@users.noreply.github.com>
Date: Thu, 9 Jan 2025 12:24:08 +0100
Subject: [PATCH] Improve test resource cleanup.

---
 .../utils/nativefst/ImmutableFSTTest.java     | 38 ++++++++++++-------
 1 file changed, 24 insertions(+), 14 deletions(-)

diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/utils/nativefst/ImmutableFSTTest.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/utils/nativefst/ImmutableFSTTest.java
index 7202874fba14..92651c51114d 100644
--- a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/utils/nativefst/ImmutableFSTTest.java
+++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/utils/nativefst/ImmutableFSTTest.java
@@ -64,8 +64,7 @@ public static void walkNode(byte[] buffer, int depth, FST fst, int node, int cnt
     }
   }
 
-  private static void verifyContent(FST fst, List<String> expected)
-      throws IOException {
+  private static void verifyContent(FST fst, List<String> expected) {
     List<String> actual = new ArrayList<>();
     for (ByteBuffer bb : fst.getSequences()) {
       assertEquals(0, bb.arrayOffset());
@@ -74,26 +73,31 @@ private static void verifyContent(FST fst, List<String> expected)
     }
     actual.sort(null);
     assertEquals(actual, expected);
-    close(fst);
   }
 
   @Test
   public void testVersion5()
       throws IOException {
+    FST fst = null;
     try (InputStream inputStream = getClass().getClassLoader().getResourceAsStream("data/abc.native.fst")) {
-      FST fst = FST.read(inputStream);
+      fst = FST.read(inputStream);
       assertFalse(fst.getFlags().contains(FSTFlags.NUMBERS));
       verifyContent(fst, _expected);
+    } finally {
+      close(fst);
     }
   }
 
   @Test
   public void testVersion5WithNumbers()
       throws IOException {
+    FST fst = null;
     try (InputStream inputStream = getClass().getClassLoader().getResourceAsStream("data/abc-numbers.native.fst")) {
-      FST fst = FST.read(inputStream);
+      fst = FST.read(inputStream);
       assertTrue(fst.getFlags().contains(FSTFlags.NUMBERS));
       verifyContent(fst, _expected);
+    } finally {
+      close(fst);
     }
   }
 
@@ -142,16 +146,22 @@ public void testSave()
     for (String input : inputList) {
       builder.add(input.getBytes(UTF_8), 0, input.length(), 127);
     }
-    FST fst = builder.complete();
-
-    File fstFile = new File(FileUtils.getTempDirectory(), "test.native.fst");
-    fst.save(new FileOutputStream(fstFile));
-
-    try (FileInputStream inputStream = new FileInputStream(fstFile)) {
-      verifyContent(FST.read(inputStream, ImmutableFST.class, true), inputList);
+    FST fst1 = builder.complete();
+    ImmutableFST fst2 = null;
+    File fstFile = null;
+    try {
+      fstFile = new File(FileUtils.getTempDirectory(), "test.native.fst");
+      fst1.save(new FileOutputStream(fstFile));
+
+      try (FileInputStream inputStream = new FileInputStream(fstFile)) {
+        fst2 = FST.read(inputStream, ImmutableFST.class, true);
+        verifyContent(fst2, inputList);
+      }
+    } finally {
+      close(fst1);
+      close(fst2);
+      FileUtils.deleteQuietly(fstFile);
     }
-
-    FileUtils.deleteQuietly(fstFile);
   }
 
   private static void close(FST fst)