diff --git a/sdks/java/extensions/google-cloud-platform-core/build.gradle b/sdks/java/extensions/google-cloud-platform-core/build.gradle index 3c6612070adb..75e0c508727a 100644 --- a/sdks/java/extensions/google-cloud-platform-core/build.gradle +++ b/sdks/java/extensions/google-cloud-platform-core/build.gradle @@ -66,6 +66,7 @@ dependencies { implementation library.java.proto_google_cloud_secret_manager_v1 implementation library.java.protobuf_java testImplementation project(path: ":sdks:java:core", configuration: "shadowTest") + testImplementation project(path: ":runners:direct-java", configuration: "shadow") testImplementation library.java.mockito_core testRuntimeOnly library.java.slf4j_jdk14 } diff --git a/sdks/java/extensions/google-cloud-platform-core/src/test/java/org/apache/beam/sdk/extensions/gcp/transforms/GcpGroupByKeyIT.java b/sdks/java/extensions/google-cloud-platform-core/src/test/java/org/apache/beam/sdk/extensions/gcp/transforms/GcpGroupByKeyIT.java index e8d3466cf097..2fd4aeb7371b 100644 --- a/sdks/java/extensions/google-cloud-platform-core/src/test/java/org/apache/beam/sdk/extensions/gcp/transforms/GcpGroupByKeyIT.java +++ b/sdks/java/extensions/google-cloud-platform-core/src/test/java/org/apache/beam/sdk/extensions/gcp/transforms/GcpGroupByKeyIT.java @@ -28,7 +28,9 @@ import com.google.protobuf.ByteString; import java.io.IOException; import java.security.SecureRandom; +import java.util.ArrayList; import java.util.Arrays; +import java.util.Collections; import java.util.List; import org.apache.beam.sdk.Pipeline; import org.apache.beam.sdk.coders.KvCoder; @@ -41,10 +43,13 @@ import org.apache.beam.sdk.transforms.Combine; import org.apache.beam.sdk.transforms.Create; import org.apache.beam.sdk.transforms.GroupByKey; +import org.apache.beam.sdk.transforms.MapElements; import org.apache.beam.sdk.transforms.Redistribute; import org.apache.beam.sdk.transforms.Sum; import org.apache.beam.sdk.values.KV; import org.apache.beam.sdk.values.PCollection; +import org.apache.beam.sdk.values.TypeDescriptor; +import org.apache.beam.sdk.values.TypeDescriptors; import org.junit.AfterClass; import org.junit.BeforeClass; import org.junit.Rule; @@ -141,6 +146,22 @@ public static void tearDown() throws IOException { } } + private static > PCollection>> sortGroupedValues( + PCollection>> input, + TypeDescriptor keyType, + TypeDescriptor valueType) { + + return input.apply( + MapElements.into(TypeDescriptors.kvs(keyType, TypeDescriptors.lists(valueType))) + .via( + kv -> { + List values = new ArrayList<>(); + kv.getValue().forEach(values::add); + Collections.sort(values); + return KV.of(kv.getKey(), values); + })); + } + @Test public void testGroupByKeyWithValidGcpSecretOption() throws Exception { if (gcpSecretVersionName == null) { @@ -165,13 +186,17 @@ public void testGroupByKeyWithValidGcpSecretOption() throws Exception { Create.of(ungroupedPairs) .withCoder(KvCoder.of(StringUtf8Coder.of(), VarIntCoder.of()))); - PCollection>> output = input.apply(GroupByKey.create()); + PCollection>> normalizedOutput = + sortGroupedValues( + input.apply(GroupByKey.create()), + TypeDescriptors.strings(), + TypeDescriptors.integers()); - PAssert.that(output) + PAssert.that(normalizedOutput) .containsInAnyOrder( KV.of("k1", Arrays.asList(3, 4)), - KV.of("k5", Arrays.asList(Integer.MAX_VALUE, Integer.MIN_VALUE)), - KV.of("k2", Arrays.asList(66, -33)), + KV.of("k5", Arrays.asList(Integer.MIN_VALUE, Integer.MAX_VALUE)), + KV.of("k2", Arrays.asList(-33, 66)), KV.of("k3", Arrays.asList(0))); p.run(); @@ -201,13 +226,17 @@ public void testGroupByKeyWithValidGcpHsmGeneratedSecretOption() throws Exceptio Create.of(ungroupedPairs) .withCoder(KvCoder.of(StringUtf8Coder.of(), VarIntCoder.of()))); - PCollection>> output = input.apply(GroupByKey.create()); + PCollection>> normalizedOutput = + sortGroupedValues( + input.apply(GroupByKey.create()), + TypeDescriptors.strings(), + TypeDescriptors.integers()); - PAssert.that(output) + PAssert.that(normalizedOutput) .containsInAnyOrder( KV.of("k1", Arrays.asList(3, 4)), - KV.of("k5", Arrays.asList(Integer.MAX_VALUE, Integer.MIN_VALUE)), - KV.of("k2", Arrays.asList(66, -33)), + KV.of("k5", Arrays.asList(Integer.MIN_VALUE, Integer.MAX_VALUE)), + KV.of("k2", Arrays.asList(-33, 66)), KV.of("k3", Arrays.asList(0))); p.run(); @@ -240,13 +269,17 @@ public void testGroupByKeyWithExistingGcpHsmGeneratedSecretOption() throws Excep Create.of(ungroupedPairs) .withCoder(KvCoder.of(StringUtf8Coder.of(), VarIntCoder.of()))); - PCollection>> output = input.apply(GroupByKey.create()); + PCollection>> normalizedOutput = + sortGroupedValues( + input.apply(GroupByKey.create()), + TypeDescriptors.strings(), + TypeDescriptors.integers()); - PAssert.that(output) + PAssert.that(normalizedOutput) .containsInAnyOrder( KV.of("k1", Arrays.asList(3, 4)), - KV.of("k5", Arrays.asList(Integer.MAX_VALUE, Integer.MIN_VALUE)), - KV.of("k2", Arrays.asList(66, -33)), + KV.of("k5", Arrays.asList(Integer.MIN_VALUE, Integer.MAX_VALUE)), + KV.of("k2", Arrays.asList(-33, 66)), KV.of("k3", Arrays.asList(0))); p.run();