fix(device): Fix the exception type that is returned when a token can…

…cellation happens (#1955)
Azure · May 13, 2021 · c3ff4c8 · c3ff4c8
1 parent ed28d8c
commit c3ff4c8
Show file tree

Hide file tree

Showing 2 changed files with 81 additions and 44 deletions.
diff --git a/e2e/test/iothub/messaging/MessageReceiveE2ETests.cs b/e2e/test/iothub/messaging/MessageReceiveE2ETests.cs
@@ -9,6 +9,7 @@
 using System.Threading.Tasks;
 using FluentAssertions;
 using Microsoft.Azure.Devices.Client;
+using Microsoft.Azure.Devices.Client.Exceptions;
 using Microsoft.Azure.Devices.Client.Transport.Mqtt;
 using Microsoft.Azure.Devices.E2ETests.Helpers;
 using Microsoft.Azure.Devices.E2ETests.Helpers.Templates;
@@ -108,6 +109,12 @@ public async Task Message_DeviceReceiveSingleMessageWithCancellationToken_Mqtt()
             await ReceiveSingleMessageWithCancellationTokenAsync(TestDeviceType.Sasl, Client.TransportType.Mqtt_Tcp_Only).ConfigureAwait(false);
         }
 
+        [LoggedTestMethod]
+        public async Task Message_DeviceReceiveSingleMessageCancelled_Mqtt()
+        {
+            await Mqtt_ReceiveSingleMessageWithCancelledAsync(TestDeviceType.Sasl).ConfigureAwait(false);
+        }
+
         [LoggedTestMethod]
         public async Task Message_DeviceReceiveSingleMessageWithCancellationToken_MqttWs()
         {
@@ -515,6 +522,22 @@ public static async Task VerifyReceivedC2DMessageAsync(Client.TransportType tran
             Assert.IsTrue(received, $"No message received for device {deviceId} with payload={payload} in {FaultInjection.RecoveryTime}.");
         }
 
+        public static async Task Mqtt_VerifyReceivedC2dMessageCancelledAsync(DeviceClient dc)
+        {
+            try
+            {
+                using var cts = new CancellationTokenSource(s_fiveSeconds);
+                await dc.ReceiveAsync(cts.Token).ConfigureAwait(false);
+            }
+            catch (IotHubCommunicationException ex)
+                when (ex.IsTransient && ex.InnerException is OperationCanceledException)
+            {
+                return;
+            }
+
+            Assert.Fail();
+        }
+
         public static async Task VerifyReceivedC2dMessageWithCancellationTokenAsync(Client.TransportType transport, DeviceClient dc, string deviceId, string payload, string p1Value, MsTestLogger logger)
         {
             var sw = new Stopwatch();
@@ -627,6 +650,19 @@ private async Task ReceiveSingleMessageAsync(TestDeviceType type, Client.Transpo
             await serviceClient.CloseAsync().ConfigureAwait(false);
         }
 
+        private async Task Mqtt_ReceiveSingleMessageWithCancelledAsync(TestDeviceType type)
+        {
+            TestDevice testDevice = await TestDevice.GetTestDeviceAsync(Logger, s_devicePrefix, type).ConfigureAwait(false);
+            using DeviceClient deviceClient = testDevice.CreateDeviceClient(Client.TransportType.Mqtt_Tcp_Only);
+
+            await deviceClient.OpenAsync().ConfigureAwait(false);
+
+            // There is no message being sent so the device client should timeout waiting for the message.
+            await Mqtt_VerifyReceivedC2dMessageCancelledAsync(deviceClient).ConfigureAwait(false);
+
+            await deviceClient.CloseAsync().ConfigureAwait(false);
+        }
+
         private async Task ReceiveSingleMessageWithCancellationTokenAsync(TestDeviceType type, Client.TransportType transport)
         {
             TestDevice testDevice = await TestDevice.GetTestDeviceAsync(Logger, s_devicePrefix, type).ConfigureAwait(false);