From c54b343abb8097b52b4d9a9dd23fba1b6171677e Mon Sep 17 00:00:00 2001
From: JoshLove-msft <54595583+JoshLove-msft@users.noreply.github.com>
Date: Fri, 4 Aug 2023 12:57:42 -0700
Subject: [PATCH] Fix MessageLockLost event test (#38021)

* Fix MessageLockLost event test

* wrap in try/finally
---
 .../tests/Processor/ProcessorLiveTests.cs     | 35 ++++++++++++++-----
 1 file changed, 26 insertions(+), 9 deletions(-)

diff --git a/sdk/servicebus/Azure.Messaging.ServiceBus/tests/Processor/ProcessorLiveTests.cs b/sdk/servicebus/Azure.Messaging.ServiceBus/tests/Processor/ProcessorLiveTests.cs
index 7d0de86c60b4a..dca8fae14c7c0 100644
--- a/sdk/servicebus/Azure.Messaging.ServiceBus/tests/Processor/ProcessorLiveTests.cs
+++ b/sdk/servicebus/Azure.Messaging.ServiceBus/tests/Processor/ProcessorLiveTests.cs
@@ -1491,10 +1491,16 @@ async Task ProcessMessage(ProcessMessageEventArgs args)
                     };
                     await args.CompleteMessageAsync(args.Message);
                     await Task.Delay(lockDuration.Add(lockDuration));
-                    Assert.IsTrue(messageLockLostRaised);
-                    Assert.IsTrue(args.MessageLockCancellationToken.IsCancellationRequested);
-                    Assert.IsFalse(args.CancellationToken.IsCancellationRequested);
-                    tcs.SetResult(true);
+                    try
+                    {
+                        Assert.IsTrue(messageLockLostRaised);
+                        Assert.IsTrue(args.MessageLockCancellationToken.IsCancellationRequested);
+                        Assert.IsFalse(args.CancellationToken.IsCancellationRequested);
+                    }
+                    finally
+                    {
+                        tcs.SetResult(true);
+                    }
                 }
                 processor.ProcessMessageAsync += ProcessMessage;
                 processor.ProcessErrorAsync += ServiceBusTestUtilities.ExceptionHandler;
@@ -1505,8 +1511,13 @@ async Task ProcessMessage(ProcessMessageEventArgs args)
             }
         }
 
+        /// <summary>
+        /// Because the message lock renewal occurs on the mgmt link, even when the connection drops, message lock renewal continues
+        /// successfully. This is in contrast to session messages where the lock renewal requires the session to be locked,
+        /// so when the connection drops, the session is lost and the lock renewal fails.
+        /// </summary>
         [Test]
-        public async Task MessageLockLostEventRaisedAfterConnectionDropped()
+        public async Task MessageLockLostEventNotRaisedAfterConnectionDropped()
         {
             var lockDuration = TimeSpan.FromSeconds(30);
             await using (var scope = await ServiceBusScope.CreateWithQueue(enablePartitioning: false, enableSession: false, lockDuration: lockDuration))
@@ -1534,10 +1545,16 @@ async Task ProcessMessage(ProcessMessageEventArgs args)
                     };
                     SimulateNetworkFailure(client);
                     await Task.Delay(lockDuration.Add(lockDuration));
-                    Assert.IsTrue(messageLockLostRaised);
-                    Assert.IsTrue(args.MessageLockCancellationToken.IsCancellationRequested);
-                    Assert.IsFalse(args.CancellationToken.IsCancellationRequested);
-                    tcs.SetResult(true);
+                    try
+                    {
+                        Assert.IsFalse(messageLockLostRaised);
+                        Assert.IsFalse(args.MessageLockCancellationToken.IsCancellationRequested);
+                        Assert.IsFalse(args.CancellationToken.IsCancellationRequested);
+                    }
+                    finally
+                    {
+                        tcs.SetResult(true);
+                    }
                 }
                 processor.ProcessMessageAsync += ProcessMessage;
                 processor.ProcessErrorAsync += ServiceBusTestUtilities.ExceptionHandler;