fix: PubSubPublishCallback handling of success and failure callbacks (#…

…1800) Fixes #1799.
GoogleCloudPlatform · May 6, 2023 · b134c92 · b134c92
1 parent 1b04165
commit b134c92
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 2 deletions.
diff --git a/...c/main/java/com/google/cloud/spring/pubsub/integration/outbound/PubSubMessageHandler.java b/...c/main/java/com/google/cloud/spring/pubsub/integration/outbound/PubSubMessageHandler.java
@@ -294,15 +294,25 @@ private class PubSubPublishCallback implements BiConsumer<String, Throwable> {
       this.message = message;
     }
 
-    @Override
-    public void accept(String messageId, Throwable throwable) {
+    private void handleSuccess(String messageId) {
       if (PubSubMessageHandler.this.successCallback != null) {
         PubSubMessageHandler.this.successCallback.onSuccess(messageId, message);
       }
+    }
 
+    private void handleFailure(Throwable throwable) {
       if (PubSubMessageHandler.this.failureCallback != null) {
         PubSubMessageHandler.this.failureCallback.onFailure(throwable, message);
       }
     }
+
+    @Override
+    public void accept(String messageId, Throwable throwable) {
+      if (throwable == null) {
+        handleSuccess(messageId);
+      } else {
+        handleFailure(throwable);
+      }
+    }
   }
 }
diff --git a/...t/java/com/google/cloud/spring/pubsub/integration/outbound/PubSubMessageHandlerTests.java b/...t/java/com/google/cloud/spring/pubsub/integration/outbound/PubSubMessageHandlerTests.java
@@ -224,18 +224,25 @@ void publishWithSuccessCallback() {
 
     AtomicReference<String> messageIdRef = new AtomicReference<>();
     AtomicReference<String> ackIdRef = new AtomicReference<>();
+    AtomicReference<Throwable> failureCauseRef = new AtomicReference<>();
 
     this.adapter.setSuccessCallback(
         (ackId, message) -> {
           messageIdRef.set(message.getHeaders().get("message_id", String.class));
           ackIdRef.set(ackId);
         });
 
+    this.adapter.setFailureCallback(
+            (exception, message) -> {
+              failureCauseRef.set(exception);
+            });
+
     this.adapter.handleMessage(testMessage);
     Awaitility.await().atMost(Duration.ofSeconds(1)).untilAtomic(messageIdRef, notNullValue());
 
     assertThat(messageIdRef).hasValue("123");
     assertThat(ackIdRef).hasValue("published12345");
+    assertThat(failureCauseRef).hasValue(null);
   }
 
   @Test
@@ -249,9 +256,15 @@ void publishWithFailureCallback() {
     Message<String> testMessage =
         new GenericMessage<>("testPayload", Collections.singletonMap("message_id", "123"));
 
+    AtomicReference<String> ackIdRef = new AtomicReference<>();
     AtomicReference<Throwable> failureCauseRef = new AtomicReference<>();
     AtomicReference<String> messageIdRef = new AtomicReference<>();
 
+    this.adapter.setSuccessCallback(
+            (ackId, message) -> {
+              ackIdRef.set(ackId);
+            });
+
     this.adapter.setFailureCallback(
         (exception, message) -> {
           failureCauseRef.set(exception);
@@ -264,5 +277,7 @@ void publishWithFailureCallback() {
     assertThat(messageIdRef).hasValue("123");
     Throwable cause = failureCauseRef.get();
     assertThat(cause).isInstanceOf(RuntimeException.class).hasMessage("boom!");
+
+    assertThat(ackIdRef).hasValue(null);
   }
 }