Fix flaky testMultiFileTableWithNaNValue test

trinodb · Mar 15, 2022 · 4e62072 · 4e62072
1 parent eb61a7a
commit 4e62072
Show file tree

Hide file tree

Showing 2 changed files with 32 additions and 19 deletions.
diff --git a/core/trino-main/src/main/java/io/trino/testing/assertions/Assert.java b/core/trino-main/src/main/java/io/trino/testing/assertions/Assert.java
@@ -56,17 +56,20 @@ public static void assertEquals(Iterable<?> actual, Iterable<?> expected, String
         }
     }
 
-    public static void assertEventually(Runnable assertion)
+    public static <E extends Exception> void assertEventually(CheckedRunnable<E> assertion)
+            throws E
     {
         assertEventually(new Duration(30, SECONDS), assertion);
     }
 
-    public static void assertEventually(Duration timeout, Runnable assertion)
+    public static <E extends Exception> void assertEventually(Duration timeout, CheckedRunnable<E> assertion)
+            throws E
     {
         assertEventually(timeout, new Duration(50, MILLISECONDS), assertion);
     }
 
-    public static void assertEventually(Duration timeout, Duration retryFrequency, Runnable assertion)
+    public static <E extends Exception> void assertEventually(Duration timeout, Duration retryFrequency, CheckedRunnable<E> assertion)
+            throws E
     {
         long start = System.nanoTime();
         while (!Thread.currentThread().isInterrupted()) {
@@ -88,4 +91,10 @@ public static void assertEventually(Duration timeout, Duration retryFrequency, R
             }
         }
     }
+
+    public interface CheckedRunnable<E extends Exception>
+    {
+        void run()
+                throws E;
+    }
 }
diff --git a/...e/src/test/java/io/trino/plugin/deltalake/AbstractTestDeltaLakeCreateTableStatistics.java b/...e/src/test/java/io/trino/plugin/deltalake/AbstractTestDeltaLakeCreateTableStatistics.java
@@ -62,6 +62,7 @@
 import static io.trino.spi.type.TimestampWithTimeZoneType.TIMESTAMP_TZ_MILLIS;
 import static io.trino.spi.type.VarcharType.createUnboundedVarcharType;
 import static io.trino.testing.TestingConnectorSession.SESSION;
+import static io.trino.testing.assertions.Assert.assertEventually;
 import static io.trino.testing.sql.TestTable.randomTableSuffix;
 import static java.lang.Double.NEGATIVE_INFINITY;
 import static java.lang.Double.POSITIVE_INFINITY;
@@ -463,22 +464,25 @@ public void testMultiFileTable()
     public void testMultiFileTableWithNaNValue()
             throws IOException
     {
-        String columnName = "orderkey";
-        DeltaLakeColumnHandle columnHandle = new DeltaLakeColumnHandle(columnName, DoubleType.DOUBLE, REGULAR);
-        try (TestTable table = new TestTable(
-                "test_partitioned_table_",
-                ImmutableList.of(columnName),
-                "SELECT IF(orderkey = 50597, nan(), CAST(orderkey AS double)) FROM tpch.tiny.orders")) {
-            List<AddFileEntry> addFileEntries = getAddFileEntries(table.getName());
-            assertThat(addFileEntries.size()).isGreaterThan(1);
-
-            List<DeltaLakeFileStatistics> statistics = addFileEntries.stream().map(entry -> entry.getStats().get()).collect(toImmutableList());
-
-            assertEquals(statistics.stream().filter(stat -> stat.getMinColumnValue(columnHandle).isEmpty() && stat.getMaxColumnValue(columnHandle).isEmpty()).count(), 1);
-            assertEquals(
-                    statistics.stream().filter(stat -> stat.getMinColumnValue(columnHandle).isPresent() && stat.getMaxColumnValue(columnHandle).isPresent()).count(),
-                    statistics.size() - 1);
-        }
+        // assertEventually because sometimes write from tpch.tiny.orders creates one file only and the test requires at least two files
+        assertEventually(() -> {
+            String columnName = "orderkey";
+            DeltaLakeColumnHandle columnHandle = new DeltaLakeColumnHandle(columnName, DoubleType.DOUBLE, REGULAR);
+            try (TestTable table = new TestTable(
+                    "test_partitioned_table_",
+                    ImmutableList.of(columnName),
+                    "SELECT IF(orderkey = 50597, nan(), CAST(orderkey AS double)) FROM tpch.tiny.orders")) {
+                List<AddFileEntry> addFileEntries = getAddFileEntries(table.getName());
+                assertThat(addFileEntries.size()).isGreaterThan(1);
+
+                List<DeltaLakeFileStatistics> statistics = addFileEntries.stream().map(entry -> entry.getStats().get()).collect(toImmutableList());
+
+                assertEquals(statistics.stream().filter(stat -> stat.getMinColumnValue(columnHandle).isEmpty() && stat.getMaxColumnValue(columnHandle).isEmpty()).count(), 1);
+                assertEquals(
+                        statistics.stream().filter(stat -> stat.getMinColumnValue(columnHandle).isPresent() && stat.getMaxColumnValue(columnHandle).isPresent()).count(),
+                        statistics.size() - 1);
+            }
+        });
     }
 
     protected class TestTable