arunpandianp commented on code in PR #38768:
URL: https://github.com/apache/beam/pull/38768#discussion_r3636569497
##########
runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/GrpcCommitWorkStreamTest.java:
##########
@@ -1134,6 +1134,76 @@ public void
testCommitWorkItem_multiplePhysicalStreams_multipleHandovers_halfClo
assertTrue(commitWorkStream.awaitTermination(10, TimeUnit.SECONDS));
}
+ @Test
+ public void testCommit_multiKeyCommit() throws Exception {
+ GrpcCommitWorkStream commitWorkStream = createCommitWorkStream();
+ FakeWindmillGrpcService.CommitStreamInfo streamInfo =
waitForConnectionAndConsumeHeader();
+
+ CompletableFuture<Windmill.CommitStatus> commitStatusFuture = new
CompletableFuture<>();
+
+ // 1. Construct two individual WorkItemCommitRequests
+ long shardingKey1 = 101L;
+ long workToken1 = 201L;
+ long cacheToken1 = 301L;
+ long shardingKey2 = 102L;
+ long workToken2 = 202L;
+ long cacheToken2 = 302L;
+ Windmill.WorkItemCommitRequest request1 =
+ Windmill.WorkItemCommitRequest.newBuilder()
+ .setKey(ByteString.copyFromUtf8("key1"))
+ .setShardingKey(shardingKey1)
+ .setWorkToken(workToken1)
+ .setCacheToken(cacheToken1)
+ .build();
+ Windmill.WorkItemCommitRequest request2 =
+ Windmill.WorkItemCommitRequest.newBuilder()
+ .setKey(ByteString.copyFromUtf8("key2"))
+ .setShardingKey(shardingKey2)
+ .setWorkToken(workToken2)
+ .setCacheToken(cacheToken2)
+ .build();
+
+ // 2. Wrap them into a MultiKeyWorkItemCommitRequest
+ Windmill.MultiKeyWorkItemCommitRequest multiKeyRequest =
+ Windmill.MultiKeyWorkItemCommitRequest.newBuilder()
+ .addRequests(request1)
+ .addRequests(request2)
+ .build();
+
+ // 3. Commit the multi-key work item using the request batcher
+ try (WindmillStream.CommitWorkStream.RequestBatcher batcher =
commitWorkStream.batcher()) {
+ assertTrue(
+ batcher.commitMultiKeyWorkItem(
+ COMPUTATION_ID, multiKeyRequest, commitStatusFuture::complete));
+ }
+
+ // 4. Receive and assert request properties on FakeWindmillGrpcService
+ Windmill.StreamingCommitWorkRequest request = streamInfo.requests.take();
+ assertThat(request.getCommitChunkCount()).isEqualTo(1);
+
+ Windmill.StreamingCommitRequestChunk chunk = request.getCommitChunk(0);
+
+ // Assert that the commit type is correctly identified as
COMMIT_TYPE_MULTI_KEY
+ assertThat(chunk.getCommitType())
+
.isEqualTo(Windmill.StreamingCommitRequestChunk.CommitType.COMMIT_TYPE_MULTI_KEY);
+
+ // Assert that the routing sharding key is mapped to the first request's
sharding key
+ assertThat(chunk.getShardingKey()).isEqualTo(request1.getShardingKey());
+
+ // Assert that the serialized payload matches the input multiKeyRequest
+ Windmill.MultiKeyWorkItemCommitRequest parsedRequest =
+
Windmill.MultiKeyWorkItemCommitRequest.parseFrom(chunk.getSerializedWorkItemCommit());
+ assertThat(parsedRequest).isEqualTo(multiKeyRequest);
+
+ // 5. Respond with the generated requestId to complete the commit
+ long requestId = chunk.getRequestId();
+ streamInfo.responseObserver.onNext(
+
Windmill.StreamingCommitResponse.newBuilder().addRequestId(requestId).build());
+
+ // 6. Verify callback completed successfully with CommitStatus.OK
+ assertThat(commitStatusFuture.get()).isEqualTo(Windmill.CommitStatus.OK);
Review Comment:
done.
--
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.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
