Sb track2 Fixing some live test (Azure#17227)

Fixing live test and also fixed the bug where receiver.complete was not completing for session enabled client

Author: hemanttanwar

sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/ServiceBusReceiverAsyncClient.java

@@ -1061,6 +1061,10 @@ deadLetterErrorDescription, propertiesToModify, sessionId, getLinkName(sessionId
             });
     }
 
+    Mono<ServiceBusAsyncConsumer> createConsumerWithReceiveLink() {
+        return Mono.fromCallable(this::getOrCreateConsumer);
+    }
+
     private ServiceBusAsyncConsumer getOrCreateConsumer() {
         final ServiceBusAsyncConsumer existing = consumer.get();
         if (existing != null) {

sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/ServiceBusSessionReceiverAsyncClient.java

@@ -113,13 +113,13 @@ public Mono<ServiceBusReceiverAsyncClient> acceptSession(String sessionId) {
         final ReceiverOptions newReceiverOptions = new ReceiverOptions(receiverOptions.getReceiveMode(),
             receiverOptions.getPrefetchCount(), receiverOptions.getMaxLockRenewDuration(),
             receiverOptions.isEnableAutoComplete(), sessionId, null);
-        final ServiceBusSessionManager sessionSpecificManager = new ServiceBusSessionManager(entityPath, entityType,
-            connectionProcessor, tracerProvider, messageSerializer, newReceiverOptions);
 
-        return sessionSpecificManager.getActiveLink().map(receiveLink -> new ServiceBusReceiverAsyncClient(
+        ServiceBusReceiverAsyncClient sessionSpecificAsyncClient = new ServiceBusReceiverAsyncClient(
             fullyQualifiedNamespace, entityPath, entityType, newReceiverOptions, connectionProcessor,
-            ServiceBusConstants.OPERATION_TIMEOUT, tracerProvider, messageSerializer, () -> { },
-            sessionSpecificManager));
+            ServiceBusConstants.OPERATION_TIMEOUT, tracerProvider, messageSerializer, () -> { });
+
+        return sessionSpecificAsyncClient.createConsumerWithReceiveLink()
+            .thenReturn(sessionSpecificAsyncClient);
     }
 
     @Override

sdk/servicebus/azure-messaging-servicebus/src/test/java/com/azure/messaging/servicebus/ServiceBusReceiverAsyncClientIntegrationTest.java

@@ -14,7 +14,6 @@
 import com.azure.messaging.servicebus.models.CompleteOptions;
 import com.azure.messaging.servicebus.models.DeadLetterOptions;
 import com.azure.messaging.servicebus.models.DeferOptions;
-import com.azure.messaging.servicebus.models.ReceiveMode;
 import com.azure.messaging.servicebus.models.SubQueue;
 import org.junit.jupiter.api.Assertions;
 import org.junit.jupiter.api.Disabled;
@@ -58,18 +57,11 @@
 class ServiceBusReceiverAsyncClientIntegrationTest extends IntegrationTestBase {
     private final ClientLogger logger = new ClientLogger(ServiceBusReceiverAsyncClientIntegrationTest.class);
     private final AtomicInteger messagesPending = new AtomicInteger();
-    private final List<Long> messagesDeferredPending = new ArrayList<>();
     private final boolean isSessionEnabled = false;
 
     private ServiceBusReceiverAsyncClient receiver;
     private ServiceBusSenderAsyncClient sender;
 
-    /**
-     * Receiver used to clean up resources in {@link #afterTest()}.
-     */
-    private ServiceBusReceiverAsyncClient receiveAndDeleteReceiver;
-    private Mono<ServiceBusReceiverAsyncClient> receiveAndDeleteReceiverMono;
-
     ServiceBusReceiverAsyncClientIntegrationTest() {
         super(new ClientLogger(ServiceBusReceiverAsyncClientIntegrationTest.class));
     }
@@ -82,52 +74,11 @@ protected void beforeTest() {
     @Override
     protected void afterTest() {
         sharedBuilder = null;
-        final int pending = messagesPending.get();
-        final int pendingDeferred = messagesDeferredPending.size();
-        if (pending < 1 && pendingDeferred < 1) {
-            dispose(receiver, sender, receiveAndDeleteReceiver);
-            if (receiveAndDeleteReceiverMono != null) {
-                dispose(receiveAndDeleteReceiverMono.block());
-            }
-            return;
-        }
-
-        // In the case that this test failed... we're going to drain the queue or subscription.
         try {
             dispose(receiver, sender);
-            if (pending > 0) {
-                if (receiveAndDeleteReceiverMono != null) {
-                    receiveAndDeleteReceiver = receiveAndDeleteReceiverMono.block();
-                }
-                receiveAndDeleteReceiver.receiveMessages()
-                    .map(message -> {
-                        logger.info("Message received: {}", message.getSequenceNumber());
-                        return message;
-                    })
-                    .timeout(Duration.ofSeconds(15), Mono.empty())
-                    .blockLast();
-            }
         } catch (Exception e) {
             logger.warning("Error occurred when draining queue.", e);
         }
-
-        try {
-            if (pendingDeferred > 0) {
-                for (Long sequenceNumber : messagesDeferredPending) {
-                    receiveAndDeleteReceiver.receiveDeferredMessage(sequenceNumber)
-                        .map(message -> {
-                            logger.info("Message received: {}", message.getSequenceNumber());
-                            return message;
-                        })
-                        .timeout(Duration.ofSeconds(15), Mono.empty())
-                        .block();
-                }
-            }
-        } catch (Exception e) {
-            logger.warning("Error occurred when draining for deferred messages.", e);
-        } finally {
-            dispose(receiveAndDeleteReceiver);
-        }
     }
 
     /**
@@ -243,9 +194,6 @@ void transactionSendReceiveAndCommit(DispositionStatus dispositionStatus) {
         StepVerifier.create(receiver.commitTransaction(transaction.get()))
             .verifyComplete();
 
-        if (dispositionStatus == DispositionStatus.DEFERRED) {
-            messagesDeferredPending.add(receivedMessage.getSequenceNumber());
-        }
     }
 
     /**
@@ -264,9 +212,6 @@ void transactionReceiveCompleteCommitMixClient(MessagingEntityType entityType) {
             .buildAsyncClient();
         this.receiver = getReceiverBuilder(useCredentials, entityType, entityIndex, shareConnection)
             .buildAsyncClient();
-        this.receiveAndDeleteReceiver = getReceiverBuilder(useCredentials, entityType, entityIndex, shareConnection)
-            .receiveMode(ReceiveMode.RECEIVE_AND_DELETE)
-            .buildAsyncClient();
 
         final String messageId = UUID.randomUUID().toString();
         final ServiceBusMessage message = getMessage(messageId, isSessionEnabled);
@@ -312,10 +257,6 @@ void receiveTwoMessagesAutoComplete(MessagingEntityType entityType, boolean isSe
             assertNotNull(sessionId, "'sessionId' should have been set.");
             this.receiver = getSessionReceiverBuilder(useCredentials, entityType, entityIndex, shareConnection)
                 .buildAsyncClient().acceptSession(sessionId).block();
-            this.receiveAndDeleteReceiverMono = getSessionReceiverBuilder(useCredentials, entityType, entityIndex,
-                shareConnection)
-                .receiveMode(ReceiveMode.RECEIVE_AND_DELETE)
-                .buildAsyncClient().acceptSession(sessionId);
         } else {
             this.receiver = getReceiverBuilder(useCredentials, entityType, entityIndex, shareConnection)
                 .buildAsyncClient();
@@ -355,16 +296,9 @@ void receiveMessageAutoComplete(MessagingEntityType entityType, boolean isSessio
             assertNotNull(sessionId, "'sessionId' should have been set.");
             this.receiver = getSessionReceiverBuilder(useCredentials, entityType, entityIndex, shareConnection)
                 .buildAsyncClient().acceptSession(sessionId).block();
-            this.receiveAndDeleteReceiverMono = getSessionReceiverBuilder(useCredentials, entityType, entityIndex,
-                shareConnection)
-                .receiveMode(ReceiveMode.RECEIVE_AND_DELETE)
-                .buildAsyncClient().acceptSession(sessionId);
         } else {
             this.receiver = getReceiverBuilder(useCredentials, entityType, entityIndex, shareConnection)
                 .buildAsyncClient();
-            this.receiveAndDeleteReceiver = getReceiverBuilder(useCredentials, entityType, entityIndex, shareConnection)
-                .receiveMode(ReceiveMode.RECEIVE_AND_DELETE)
-                .buildAsyncClient();
         }
 
         final String messageId = UUID.randomUUID().toString();
@@ -402,6 +336,12 @@ void peekMessage(MessagingEntityType entityType, boolean isSessionEnabled) {
         StepVerifier.create(receiver.peekMessage())
             .assertNext(receivedMessage -> assertMessageEquals(receivedMessage, messageId, isSessionEnabled))
             .verifyComplete();
+
+        // cleanup
+        StepVerifier.create(receiver.receiveMessages())
+            .assertNext(receivedMessage -> receiver.complete(receivedMessage).block(OPERATION_TIMEOUT))
+            .thenCancel()
+            .verify();
     }
 
     /**
@@ -436,9 +376,10 @@ void sendScheduledMessageAndReceive(MessagingEntityType entityType, boolean isSe
         sender.scheduleMessage(message, scheduledEnqueueTime).block(TIMEOUT);
 
         // Assert & Act
-        StepVerifier.create(Mono.delay(Duration.ofSeconds(3)).then(receiveAndDeleteReceiver.receiveMessages().next()))
+        StepVerifier.create(Mono.delay(Duration.ofSeconds(4)).then(receiver.receiveMessages().next()))
             .assertNext(receivedMessage -> {
                 assertMessageEquals(receivedMessage, messageId, isSessionEnabled);
+                receiver.complete(receivedMessage).block(OPERATION_TIMEOUT);
                 messagesPending.decrementAndGet();
             })
             .verifyComplete();
@@ -558,9 +499,15 @@ void peekMessages(MessagingEntityType entityType, boolean isSessionEnabled) {
                 .assertNext(message -> checkCorrectMessage.accept(message, 7))
                 .verifyComplete();
         } finally {
-            receiveAndDeleteReceiver.receiveMessages()
-                .take(messages.size())
-                .blockLast(Duration.ofSeconds(15));
+            AtomicInteger completed = new AtomicInteger();
+            StepVerifier.create(receiver.receiveMessages().take(messages.size()))
+                .thenConsumeWhile(receivedMessage -> {
+                    completed.incrementAndGet();
+                    receiver.complete(receivedMessage).block(OPERATION_TIMEOUT);
+                    return completed.get() <= messages.size();
+                })
+                .expectComplete()
+                .verify();
 
             messagesPending.addAndGet(-messages.size());
         }
@@ -586,6 +533,17 @@ void peekMessagesFromSequence(MessagingEntityType entityType) {
         StepVerifier.create(receiver.peekMessagesAt(maxMessages, fromSequenceNumber))
             .expectNextCount(maxMessages)
             .verifyComplete();
+
+        // cleanup
+        StepVerifier.create(receiver.receiveMessages().take(maxMessages))
+            .assertNext(receivedMessage -> {
+                receiver.complete(receivedMessage).block(Duration.ofSeconds(15));
+            })
+            .assertNext(receivedMessage -> {
+                receiver.complete(receivedMessage).block(Duration.ofSeconds(15));
+            })
+            .expectComplete()
+            .verify(TIMEOUT);
     }
 
     /**
@@ -795,7 +753,7 @@ void receiveAndDefer(MessagingEntityType entityType, boolean isSessionEnabled) {
     @ParameterizedTest
     void receiveDeferredMessageBySequenceNumber(MessagingEntityType entityType, DispositionStatus dispositionStatus) {
         // Arrange
-        setSenderAndReceiver(entityType, 0, false);
+        setSenderAndReceiver(entityType, TestUtils.USE_CASE_DEFERRED_MESSAGE_BY_SEQUENCE_NUMBER, false);
 
         final String messageId = UUID.randomUUID().toString();
         final ServiceBusMessage message = getMessage(messageId, false);
@@ -817,7 +775,6 @@ void receiveDeferredMessageBySequenceNumber(MessagingEntityType entityType, Disp
         switch (dispositionStatus) {
             case ABANDONED:
                 operation = receiver.abandon(receivedDeferredMessage);
-                messagesDeferredPending.add(receivedDeferredMessage.getSequenceNumber());
                 break;
             case SUSPENDED:
                 operation = receiver.deadLetter(receivedDeferredMessage);
@@ -844,9 +801,9 @@ void receiveDeferredMessageBySequenceNumber(MessagingEntityType entityType, Disp
     @ParameterizedTest
     void sendReceiveMessageWithVariousPropertyTypes(MessagingEntityType entityType) {
         // Arrange
+        final boolean isSessionEnabled = true;
         setSenderAndReceiver(entityType, TestUtils.USE_CASE_SEND_RECEIVE_WITH_PROPERTIES, isSessionEnabled);
 
-        final boolean isSessionEnabled = true;
         final String messageId = UUID.randomUUID().toString();
         final ServiceBusMessage messageToSend = getMessage(messageId, isSessionEnabled);
 
@@ -866,7 +823,7 @@ void sendReceiveMessageWithVariousPropertyTypes(MessagingEntityType entityType)
         sendMessage(messageToSend).block(TIMEOUT);
 
         // Assert & Act
-        StepVerifier.create(receiveAndDeleteReceiver.receiveMessages())
+        StepVerifier.create(receiver.receiveMessages())
             .assertNext(receivedMessage -> {
                 messagesPending.decrementAndGet();
                 assertMessageEquals(receivedMessage, messageId, isSessionEnabled);
@@ -887,6 +844,7 @@ void sendReceiveMessageWithVariousPropertyTypes(MessagingEntityType entityType)
                             actual));
                     }
                 }
+                receiver.complete(receivedMessage).block(OPERATION_TIMEOUT);
             })
             .thenCancel()
             .verify();
@@ -1232,19 +1190,11 @@ private void setSenderAndReceiver(MessagingEntityType entityType, int entityInde
             this.receiver = getSessionReceiverBuilder(useCredentials, entityType, entityIndex, shareConnection)
                 .disableAutoComplete()
                 .buildAsyncClient().acceptSession(sessionId).block();
-            this.receiveAndDeleteReceiverMono = getSessionReceiverBuilder(useCredentials, entityType, entityIndex,
-                shareConnection)
-                .disableAutoComplete()
-                .receiveMode(ReceiveMode.RECEIVE_AND_DELETE)
-                .buildAsyncClient().acceptSession(sessionId);
+
         } else {
             this.receiver = getReceiverBuilder(useCredentials, entityType, entityIndex, shareConnection)
                 .disableAutoComplete()
                 .buildAsyncClient();
-            this.receiveAndDeleteReceiver = getReceiverBuilder(useCredentials, entityType, entityIndex, shareConnection)
-                .disableAutoComplete()
-                .receiveMode(ReceiveMode.RECEIVE_AND_DELETE)
-                .buildAsyncClient();
         }
     }
 

sdk/servicebus/azure-messaging-servicebus/src/test/java/com/azure/messaging/servicebus/TestUtils.java

@@ -51,7 +51,7 @@ public class TestUtils {
     static final int USE_CASE_MULTIPLE_RECEIVE_ONE_TIMEOUT = 5;
     static final int USE_CASE_PEEK_BATCH_MESSAGES = 6;
     static final int USE_CASE_SEND_READ_BACK_MESSAGES = 7;
-    static final int USE_CASE_MULTIPLE_SESSION = 8;
+    static final int USE_CASE_DEFERRED_MESSAGE_BY_SEQUENCE_NUMBER = 8;
     static final int USE_CASE_PEEK_MESSAGE_FROM_SEQUENCE = 9;
     static final int USE_CASE_PEEK_RECEIVE_AND_DEFER = 10;
     static final int USE_CASE_PEEK_TRANSACTION_SENDRECEIVE_AND_COMPLETE = 11;