update FileRoundtrip validate method with unit test

llama90 · Jun 11, 2024 · 2065fcf · 2065fcf
1 parent 4df00fa
commit 2065fcf
Show file tree

Hide file tree

Showing 2 changed files with 53 additions and 3 deletions.
diff --git a/java/tools/src/main/java/org/apache/arrow/tools/FileRoundtrip.java b/java/tools/src/main/java/org/apache/arrow/tools/FileRoundtrip.java
@@ -52,13 +52,29 @@ public static void main(String[] args) {
     System.exit(new FileRoundtrip(System.err).run(args));
   }
 
-  private File validateFile(String type, String fileName) {
+  private File validateFile(String type, String fileName) throws IOException {
     if (fileName == null) {
       throw new IllegalArgumentException("missing " + type + " file parameter");
     }
     File f = new File(fileName);
-    if (!f.exists() || f.isDirectory()) {
-      throw new IllegalArgumentException(type + " file not found: " + f.getAbsolutePath());
+    if (type.equals("input")) {
+      if (!f.exists() || f.isDirectory()) {
+        throw new IllegalArgumentException(type + " file not found: " + f.getAbsolutePath());
+      }
+    } else if (type.equals("output")) {
+      File parentDir = f.getParentFile();
+      if (parentDir != null && !parentDir.exists()) {
+        if (!parentDir.mkdirs()) {
+          throw new IOException(
+              "Failed to create parent directory: " + parentDir.getAbsolutePath());
+        }
+      }
+      if (!f.exists()) {
+        if (!f.createNewFile()) {
+          throw new IOException(
+              type + " file not found and could not be created: " + f.getAbsolutePath());
+        }
+      }
     }
     return f;
   }

diff --git a/java/tools/src/test/java/org/apache/arrow/tools/TestFileRoundtrip.java b/java/tools/src/test/java/org/apache/arrow/tools/TestFileRoundtrip.java
@@ -19,10 +19,12 @@
 import static org.apache.arrow.tools.ArrowFileTestFixtures.validateOutput;
 import static org.apache.arrow.tools.ArrowFileTestFixtures.writeInput;
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertThrows;
 
 import java.io.File;
 import org.apache.arrow.memory.BufferAllocator;
 import org.apache.arrow.memory.RootAllocator;
+import org.apache.arrow.vector.ipc.InvalidArrowFileException;
 import org.junit.After;
 import org.junit.Before;
 import org.junit.Rule;
@@ -32,6 +34,7 @@
 public class TestFileRoundtrip {
 
   @Rule public TemporaryFolder testFolder = new TemporaryFolder();
+  @Rule public TemporaryFolder testAnotherFolder = new TemporaryFolder();
 
   private BufferAllocator allocator;
 
@@ -58,4 +61,35 @@ public void test() throws Exception {
 
     validateOutput(testOutFile, allocator);
   }
+
+  @Test
+  public void testDiffFolder() throws Exception {
+    File testInFile = testFolder.newFile("testIn.arrow");
+    File testOutFile = testAnotherFolder.newFile("testOut.arrow");
+
+    writeInput(testInFile, allocator);
+
+    String[] args = {"-i", testInFile.getAbsolutePath(), "-o", testOutFile.getAbsolutePath()};
+    int result = new FileRoundtrip(System.err).run(args);
+    assertEquals(0, result);
+
+    validateOutput(testOutFile, allocator);
+  }
+
+  @Test
+  public void testNotPreparedInput() throws Exception {
+    File testInFile = testFolder.newFile("testIn.arrow");
+    File testOutFile = testFolder.newFile("testOut.arrow");
+
+    String[] args = {"-i", testInFile.getAbsolutePath(), "-o", testOutFile.getAbsolutePath()};
+
+    // In JUnit 5, since the file itself is not created, the exception and message will be different.
+    Exception exception =
+        assertThrows(
+            InvalidArrowFileException.class,
+            () -> {
+              new FileRoundtrip(System.err).run(args);
+            });
+    assertEquals("file too small: 0", exception.getMessage());
+  }
 }