diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/testing/PAssert.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/testing/PAssert.java index 4a73361dfc4f..99d41ad84242 100644 --- a/sdks/java/core/src/main/java/org/apache/beam/sdk/testing/PAssert.java +++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/testing/PAssert.java @@ -233,6 +233,19 @@ public interface IterableAssert { */ IterableAssert inWindow(BoundedWindow window); + /** + * Creates a new {@link IterableAssert} like this one, but with the assertion restricted to only + * run on the provided window. + * + *

The assertion will expect outputs to be produced to the provided window exactly once. If + * the upstream {@link Trigger} may produce output multiple times, consider instead using {@link + * #inFinalPane(BoundedWindow)} or {@link #inOnTimePane(BoundedWindow)}. + * + * @return a new {@link IterableAssert} like this one but with the assertion only applied to the + * specified window. + */ + IterableAssert inOnlyPane(BoundedWindow window); + /** * Creates a new {@link IterableAssert} like this one, but with the assertion restricted to only * run on the provided window, running the checker only on the final pane for each key. @@ -340,6 +353,20 @@ public interface IterableAssert { /** Builder interface for assertions applicable to a single value. */ public interface SingletonAssert { + /** + * Creates a new {@link SingletonAssert} like this one, but with the assertion restricted to + * only run on the provided window. + * + *

The assertion will concatenate all panes present in the provided window if the {@link + * Trigger} produces multiple panes. If the windowing strategy accumulates fired panes and + * triggers fire multple times, consider using instead {@link #inFinalPane(BoundedWindow)} or + * {@link #inOnTimePane(BoundedWindow)}. + * + * @return a new {@link SingletonAssert} like this one but with the assertion only applied to + * the specified window. + */ + SingletonAssert inWindow(BoundedWindow window); + /** * Creates a new {@link SingletonAssert} like this one, but with the assertion restricted to * only run on the provided window. @@ -631,6 +658,11 @@ public PCollectionContentsAssert inWindow(BoundedWindow window) { return withPane(window, PaneExtractors.allPanes()); } + @Override + public PCollectionContentsAssert inOnlyPane(BoundedWindow window) { + return withPane(window, PaneExtractors.onlyPane(site)); + } + @Override public PCollectionContentsAssert inFinalPane(BoundedWindow window) { return withPane(window, PaneExtractors.finalPane()); @@ -822,6 +854,11 @@ public PCollectionSingletonIterableAssert inWindow(BoundedWindow window) { return withPanes(window, PaneExtractors.allPanes()); } + @Override + public PCollectionSingletonIterableAssert inOnlyPane(BoundedWindow window) { + return withPanes(window, PaneExtractors.onlyPane(site)); + } + @Override public PCollectionSingletonIterableAssert inFinalPane(BoundedWindow window) { return withPanes(window, PaneExtractors.finalPane()); @@ -943,6 +980,11 @@ private static class PCollectionSingletonAssert implements SingletonAssert this.site = site; } + @Override + public PCollectionSingletonAssert inWindow(BoundedWindow window) { + return withPanes(window, PaneExtractors.allPanes()); + } + @Override public PCollectionSingletonAssert inFinalPane(BoundedWindow window) { return withPanes(window, PaneExtractors.finalPane()); @@ -1071,6 +1113,11 @@ private PCollectionViewAssert( this.site = site; } + @Override + public PCollectionViewAssert inWindow(BoundedWindow window) { + return inPane(window, PaneExtractors.allPanes()); + } + @Override public PCollectionViewAssert inOnlyPane(BoundedWindow window) { return inPane(window, PaneExtractors.onlyPane(site));