[GitHub] [beam] KhaninArtur commented on a change in pull request #14322: [BEAM-449] Support PCollectionList in PAssert

[GitBox <gi...@apache.org>]

KhaninArtur commented on a change in pull request #14322:
URL: https://github.com/apache/beam/pull/14322#discussion_r600470178



##########
File path: sdks/java/core/src/main/java/org/apache/beam/sdk/testing/PAssert.java
##########
@@ -512,8 +513,72 @@ public int hashCode() {
         PAssertionSite.capture(reason));
   }
 
+  /**
+   * Constructs an {@link PCollectionListContentsAssert} for the provided {@link PCollectionList}.
+   */
+  public static <T> PCollectionListContentsAssert<T> thatList(PCollectionList<T> actual) {
+    return new PCollectionListContentsAssert<>(actual);
+  }
+
+  /**
+   * Constructs an {@link IterableAssert} for the elements of the flattened {@link PCollectionList}.
+   */
+  public static <T> IterableAssert<T> thatFlattened(PCollectionList<T> actual) {
+    PCollection<T> flatten = actual.apply(Flatten.pCollections());
+    return that(flatten.getName(), flatten);
+  }
+
+  /**
+   * Constructs an {@link IterableAssert} for the elements of the flattened {@link PCollectionList}
+   * with the specified reason.
+   */
+  public static <T> IterableAssert<T> thatFlattened(String reason, PCollectionList<T> actual) {
+    return new PCollectionContentsAssert<>(
+        actual.apply(Flatten.pCollections()), PAssertionSite.capture(reason));
+  }
+
   ////////////////////////////////////////////////////////////
 
+  /**
+   * An assert about the contents of each {@link PCollection} in the given {@link PCollectionList}.
+   */
+  protected static class PCollectionListContentsAssert<T> {
+    private final PCollectionList<T> pCollectionList;
+
+    public PCollectionListContentsAssert(PCollectionList<T> actual) {
+      this.pCollectionList = actual;
+    }
+
+    /**
+     * Applies one {@link SerializableFunction} to check the elements of each {@code PCollection} in
+     * the {@link #pCollectionList}.

Review comment:
       Shouldn't it be PCollectionList?

##########
File path: sdks/java/core/src/main/java/org/apache/beam/sdk/testing/PAssert.java
##########
@@ -512,8 +513,72 @@ public int hashCode() {
         PAssertionSite.capture(reason));
   }
 
+  /**
+   * Constructs an {@link PCollectionListContentsAssert} for the provided {@link PCollectionList}.
+   */
+  public static <T> PCollectionListContentsAssert<T> thatList(PCollectionList<T> actual) {
+    return new PCollectionListContentsAssert<>(actual);
+  }
+
+  /**
+   * Constructs an {@link IterableAssert} for the elements of the flattened {@link PCollectionList}.
+   */
+  public static <T> IterableAssert<T> thatFlattened(PCollectionList<T> actual) {
+    PCollection<T> flatten = actual.apply(Flatten.pCollections());
+    return that(flatten.getName(), flatten);
+  }
+
+  /**
+   * Constructs an {@link IterableAssert} for the elements of the flattened {@link PCollectionList}
+   * with the specified reason.
+   */
+  public static <T> IterableAssert<T> thatFlattened(String reason, PCollectionList<T> actual) {
+    return new PCollectionContentsAssert<>(
+        actual.apply(Flatten.pCollections()), PAssertionSite.capture(reason));
+  }
+
   ////////////////////////////////////////////////////////////
 
+  /**
+   * An assert about the contents of each {@link PCollection} in the given {@link PCollectionList}.
+   */
+  protected static class PCollectionListContentsAssert<T> {
+    private final PCollectionList<T> pCollectionList;
+
+    public PCollectionListContentsAssert(PCollectionList<T> actual) {
+      this.pCollectionList = actual;
+    }
+
+    /**
+     * Applies one {@link SerializableFunction} to check the elements of each {@code PCollection} in
+     * the {@link #pCollectionList}.
+     *
+     * <p>Returns this {@code PCollectionListContentsAssert}.
+     */
+    public PCollectionListContentsAssert<T> satisfies(
+        SerializableFunction<Iterable<T>, Void> checkerFn) {
+      for (int i = 0; i < pCollectionList.size(); i++) {
+        PAssert.that(pCollectionList.get(i)).satisfies(checkerFn);
+      }
+      return this;
+    }
+
+    /**
+     * Takes list of {@link SerializableFunction}s of the same size as {@link #pCollectionList}, and
+     * applies each matcher to the {@code PCollection} with the identical index in the {@link
+     * #pCollectionList}.
+     *
+     * <p>Returns this {@code PCollectionListContentsAssert}.
+     */
+    public PCollectionListContentsAssert<T> satisfies(
+        List<SerializableFunction<Iterable<T>, Void>> checkerFnList) {
+      for (int i = 0; i < pCollectionList.size(); i++) {
+        PAssert.that(pCollectionList.get(i)).satisfies(checkerFnList.get(i));

Review comment:
       Is there somewhere any validation that the size of the collection list and checker list are equal?




-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

For queries about this service, please contact Infrastructure at:
users@infra.apache.org