Unwrap exceptions in instrumented LROs (Azure#20540)

Author: pakrym

sdk/core/Azure.Core.TestFramework/src/OperationInterceptor.cs

@@ -4,6 +4,7 @@
 using System;
 using System.Linq;
 using System.Reflection;
+using System.Runtime.ExceptionServices;
 using System.Threading;
 using Castle.DynamicProxy;
 
@@ -31,15 +32,29 @@ public void Intercept(IInvocation invocation)
                     CheckArguments(invocation.Arguments);
                     var cancellationToken = invocation.Arguments.Last();
                     var waitForCompletionMethod = invocation.TargetType.GetMethod(WaitForCompletionMethodName, new[]{typeof(TimeSpan), typeof(CancellationToken)});
-                    invocation.ReturnValue = waitForCompletionMethod.Invoke(invocation.InvocationTarget, new[] {NoWaitDelay, cancellationToken});
+                    try
+                    {
+                        invocation.ReturnValue = waitForCompletionMethod.Invoke(invocation.InvocationTarget, new[] {NoWaitDelay, cancellationToken});
+                    }
+                    catch (TargetInvocationException ex)
+                    {
+                        ExceptionDispatchInfo.Capture(ex.InnerException).Throw();
+                    }
                     return;
                 }
 
                 if (invocation.Method.Name == WaitForCompletionResponseAsync.Name)
                 {
                     CheckArguments(invocation.Arguments);
                     var cancellationToken = invocation.Arguments.Last();
-                    invocation.ReturnValue = WaitForCompletionResponseAsync.Invoke(invocation.InvocationTarget, new[] {NoWaitDelay, cancellationToken});
+                    try
+                    {
+                        invocation.ReturnValue = WaitForCompletionResponseAsync.Invoke(invocation.InvocationTarget, new[] {NoWaitDelay, cancellationToken});
+                    }
+                    catch (TargetInvocationException ex)
+                    {
+                        ExceptionDispatchInfo.Capture(ex.InnerException).Throw();
+                    }
                     return;
                 }
             }

sdk/core/Azure.Core/tests/RecordedTestBaseTests.cs

@@ -3,7 +3,6 @@
 
 using System;
 using System.Collections.Generic;
-using System.Linq;
 using System.Threading;
 using System.Threading.Tasks;
 using Azure.Core.TestFramework;
@@ -80,21 +79,42 @@ public async Task WaitForCompletionWaitsDuringLive()
             Assert.AreEqual(TimeSpan.FromSeconds(10), original.WaitForCompletionCalls[3]);
         }
 
+        [Test]
+        public async Task WaitForCompletionErrorsArePropagated()
+        {
+            var client = InstrumentClient(new RecordedClient());
+
+            var operation = await client.StartOperationAsync(true);
+
+            Assert.ThrowsAsync<RequestFailedException>(async () => await operation.WaitForCompletionAsync());
+        }
+
         public class RecordedClient
         {
-            public virtual Task<CustomOperation> StartOperationAsync()
+            public virtual Task<CustomOperation> StartOperationAsync(bool throwOnWait = false)
             {
-                return Task.FromResult(new CustomOperation());
+                return Task.FromResult(new CustomOperation(throwOnWait));
             }
 
-            public virtual CustomOperation StartOperation()
+            public virtual CustomOperation StartOperation(bool throwOnWait = false)
             {
-                return new CustomOperation();
+                return new CustomOperation(throwOnWait);
             }
         }
 
         public class CustomOperation : Operation<int>
         {
+            private readonly bool _throwOnWait;
+
+            public CustomOperation(bool throwOnWait)
+            {
+                _throwOnWait = throwOnWait;
+            }
+
+            protected CustomOperation()
+            {
+            }
+
             public virtual List<TimeSpan?> WaitForCompletionCalls { get; } = new();
             public override string Id { get; }
             public override Response GetRawResponse()
@@ -123,6 +143,10 @@ public override ValueTask<Response<int>> WaitForCompletionAsync(CancellationToke
             public override ValueTask<Response<int>> WaitForCompletionAsync(TimeSpan pollingInterval, CancellationToken cancellationToken)
             {
                 WaitForCompletionCalls.Add(pollingInterval);
+                if (_throwOnWait)
+                {
+                    throw new RequestFailedException(400, "Operation failed");
+                }
                 return new(Response.FromValue(0, new MockResponse(200)));
             }
         }