prodriguezdefino commented on code in PR #40606:
URL: https://github.com/apache/beam/pull/40606#discussion_r4300033834
##########
sdks/java/core/src/test/java/org/apache/beam/sdk/fn/data/BeamFnDataInboundObserverTest.java:
##########
@@ -218,6 +220,84 @@ public void
testBadProducerDataFailureVisibleToAwaitCompletionCallerAndProducer(
future.get();
}
+ @Test
+ public void testAcceptWithCallbackInvokedOnProcessing() throws Exception {
+ Collection<WindowedValue<String>> values = new ArrayList<>();
+ BeamFnDataInboundObserver observer =
+ BeamFnDataInboundObserver.forConsumers(
+ Collections.singletonList(DataEndpoint.create(TRANSFORM_ID, CODER,
values::add)),
+ Collections.emptyList());
+
+ AtomicBoolean callback1Ran = new AtomicBoolean();
+ AtomicBoolean callback2Ran = new AtomicBoolean();
+ AtomicBoolean callbackLastRan = new AtomicBoolean();
+
+ Future<?> future =
+ executor.submit(
+ () -> {
+ observer.acceptWithCallback(dataWith("ABC"), () ->
callback1Ran.set(true));
+ assertTrue(
+ observer.tryAcceptWithCallback(dataWith("DEF"), () ->
callback2Ran.set(true)));
+ observer.acceptWithCallback(lastData(), () ->
callbackLastRan.set(true));
+ return null;
+ });
+
+ observer.awaitCompletion();
+ future.get();
+
+ assertTrue(callback1Ran.get());
+ assertTrue(callback2Ran.get());
+ assertTrue(callbackLastRan.get());
+ assertThat(
+ values,
+ contains(
+ WindowedValues.valueInGlobalWindow("ABC"),
WindowedValues.valueInGlobalWindow("DEF")));
+ }
+
+ @Test
+ public void testCallbacksInvokedOnClose() throws Exception {
+ BeamFnDataInboundObserver observer =
+ BeamFnDataInboundObserver.forConsumers(
+ Collections.singletonList(DataEndpoint.create(TRANSFORM_ID, CODER,
val -> {})),
+ Collections.emptyList());
+
+ AtomicBoolean callback1 = new AtomicBoolean();
+ AtomicBoolean callback2 = new AtomicBoolean();
+
+ observer.acceptWithCallback(dataWith("ABC"), () -> callback1.set(true));
+ observer.tryAcceptWithCallback(dataWith("DEF"), () -> callback2.set(true));
+
+ observer.close();
+
+ assertTrue(callback1.get());
+ assertTrue(callback2.get());
+ }
+
+ @Test
+ public void testCallbacksInvokedOnConsumerException() throws Exception {
+ BeamFnDataInboundObserver observer =
+ BeamFnDataInboundObserver.forConsumers(
+ Collections.singletonList(
+ DataEndpoint.create(
+ TRANSFORM_ID,
+ CODER,
+ val -> {
+ throw new RuntimeException("consumer error");
+ })),
+ Collections.emptyList());
+
+ AtomicBoolean callback1 = new AtomicBoolean();
+ AtomicBoolean callback2 = new AtomicBoolean();
+
+ observer.acceptWithCallback(dataWith("ABC"), () -> callback1.set(true));
+ observer.acceptWithCallback(dataWith("DEF"), () -> callback2.set(true));
+
+ assertThrows(RuntimeException.class, observer::awaitCompletion);
+
+ assertFalse(callback1.get());
Review Comment:
Agreed. I'd changed this to match the bug. It asserts
`assertTrue(callback1.get())` now.
--
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]
