diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/io/FileIOTest.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/io/FileIOTest.java index dffc4943bfab..2cffce762135 100644 --- a/sdks/java/core/src/test/java/org/apache/beam/sdk/io/FileIOTest.java +++ b/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 { + private final int numOutputs; + + public AfterNumberOfNewOutputs(int numOutputs) { + this.numOutputs = numOutputs; + } + + @Override + public org.apache.beam.sdk.coders.Coder 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,7 +329,9 @@ 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 matchAllMetadata = p.apply("create for matchAll new files", Create.of(watchPath.resolve("*").toString())) .apply( @@ -299,7 +339,9 @@ public void testMatchWatchForNewFiles() throws IOException, InterruptedException 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 matchUpdatedMetadata = p.apply( "match updated", @@ -307,7 +349,9 @@ public void testMatchWatchForNewFiles() throws IOException, InterruptedException .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 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