Merge pull request #14981 from y1chi/BEAM-12470

[BEAM-12470] Increase input size of ReshuffleTest.testAssignShardFn t…
apache · Jun 9, 2021 · 81a46a9 · 81a46a9
2 parents 803a6aa + f3b1f5e
commit 81a46a9
Showing 1 changed file with 7 additions and 1 deletion.
diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/ReshuffleTest.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/ReshuffleTest.java
@@ -26,6 +26,7 @@
 import static org.junit.Assert.assertEquals;
 
 import java.io.Serializable;
+import java.util.List;
 import org.apache.beam.sdk.coders.KvCoder;
 import org.apache.beam.sdk.coders.StringUtf8Coder;
 import org.apache.beam.sdk.coders.VarIntCoder;
@@ -47,6 +48,7 @@
 import org.apache.beam.sdk.values.TypeDescriptors;
 import org.apache.beam.sdk.values.WindowingStrategy;
 import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Lists;
 import org.joda.time.Duration;
 import org.joda.time.Instant;
 import org.junit.Rule;
@@ -280,10 +282,14 @@ public void testReshuffleWithTimestampsStreaming() {
   @Test
   @Category({ValidatesRunner.class})
   public void testAssignShardFn() {
+    List<KV<String, Integer>> inputKvs = Lists.newArrayList();
+    for (int i = 0; i < 10; i++) {
+      inputKvs.addAll(ARBITRARY_KVS);
+    }
 
     PCollection<KV<String, Integer>> input =
         pipeline.apply(
-            Create.of(ARBITRARY_KVS).withCoder(KvCoder.of(StringUtf8Coder.of(), VarIntCoder.of())));
+            Create.of(inputKvs).withCoder(KvCoder.of(StringUtf8Coder.of(), VarIntCoder.of())));
 
     PCollection<Integer> output =
         input