fix flakey fileiotest (#39174)

* fix flakey fileiotest

* fix gemini

* more gemini fixes

Author: derrickaw

sdks/java/core/src/test/java/org/apache/beam/sdk/io/FileIOTest.java

@@ -271,9 +271,47 @@ public void processElement(ProcessContext context, @StateId("count") ValueState<
     private final String watchPathStr;
   }
 
+  private static class AfterNumberOfNewOutputs
+      implements Watch.Growth.TerminationCondition<String, Integer> {
+    private final int numOutputs;
+
+    public AfterNumberOfNewOutputs(int numOutputs) {
+      this.numOutputs = numOutputs;
+    }
+
+    @Override
+    public org.apache.beam.sdk.coders.Coder<Integer> getStateCoder() {
+      return VarIntCoder.of();
+    }
+
+    @Override
+    public Integer forNewInput(org.joda.time.Instant now, String input) {
+      return 0;
+    }
+
+    @Override
+    public Integer onSeenNewOutput(org.joda.time.Instant now, Integer state) {
+      return (state == null ? 0 : state) + 1;
+    }
+
+    @Override
+    public boolean canStopPolling(org.joda.time.Instant now, Integer state) {
+      return (state == null ? 0 : state) >= numOutputs;
+    }
+
+    @Override
+    public String toString(Integer state) {
+      return "AfterNumberOfNewOutputs(" + state + "/" + numOutputs + ")";
+    }
+  }
+
   @Test
   @Category({NeedsRunner.class, UsesUnboundedSplittableParDo.class})
   public void testMatchWatchForNewFiles() throws IOException, InterruptedException {
+    final int numFiles = 3;
+    final int numFirstFiles = 3;
+    final int numUpdatedFiles = 6; // first x3, second x2, third x1
+
     // Write some files to a "source" directory.
     final Path sourcePath = tmpFolder.getRoot().toPath().resolve("source");
     sourcePath.toFile().mkdir();
@@ -291,23 +329,29 @@ public void testMatchWatchForNewFiles() throws IOException, InterruptedException
                 .filepattern(watchPath.resolve("*").toString())
                 .continuously(
                     Duration.millis(100),
-                    Watch.Growth.afterTimeSinceNewOutput(Duration.standardSeconds(1))));
+                    Watch.Growth.eitherOf(
+                        new AfterNumberOfNewOutputs(numFiles),
+                        Watch.Growth.afterTimeSinceNewOutput(Duration.standardSeconds(10)))));
     PCollection<MatchResult.Metadata> matchAllMetadata =
         p.apply("create for matchAll new files", Create.of(watchPath.resolve("*").toString()))
             .apply(
                 "match filename through matchAll",
                 FileIO.matchAll()
                     .continuously(
                         Duration.millis(100),
-                        Watch.Growth.afterTimeSinceNewOutput(Duration.standardSeconds(1))));
+                        Watch.Growth.eitherOf(
+                            new AfterNumberOfNewOutputs(numFiles),
+                            Watch.Growth.afterTimeSinceNewOutput(Duration.standardSeconds(10)))));
     PCollection<MatchResult.Metadata> matchUpdatedMetadata =
         p.apply(
             "match updated",
             FileIO.match()
                 .filepattern(watchPath.resolve("first").toString())
                 .continuously(
                     Duration.millis(100),
-                    Watch.Growth.afterTimeSinceNewOutput(Duration.standardSeconds(1)),
+                    Watch.Growth.eitherOf(
+                        new AfterNumberOfNewOutputs(numFirstFiles),
+                        Watch.Growth.afterTimeSinceNewOutput(Duration.standardSeconds(10))),
                     true));
     PCollection<MatchResult.Metadata> matchAllUpdatedMetadata =
         p.apply("create for matchAll updated files", Create.of(watchPath.resolve("*").toString()))
@@ -316,7 +360,9 @@ public void testMatchWatchForNewFiles() throws IOException, InterruptedException
                 FileIO.matchAll()
                     .continuously(
                         Duration.millis(100),
-                        Watch.Growth.afterTimeSinceNewOutput(Duration.standardSeconds(1)),
+                        Watch.Growth.eitherOf(
+                            new AfterNumberOfNewOutputs(numUpdatedFiles),
+                            Watch.Growth.afterTimeSinceNewOutput(Duration.standardSeconds(10))),
                         true));
 
     // write one file at the beginning. This will trigger the first output for matchAll