Instrument mock operations (Azure#27399)

* Instrument mock operations

Resolves Azure#27393

* Resolve PR feedback

* Resolve offline feedback

Author: heaths

common/ManagementTestShared/Redesign/ManagementRecordedTestBase.cs

@@ -8,7 +8,6 @@
 using System;
 using System.Collections.Generic;
 using System.Linq;
-using System.Reflection;
 using System.Threading.Tasks;
 
 namespace Azure.ResourceManager.TestFramework
@@ -35,6 +34,8 @@ public abstract class ManagementRecordedTestBase<TEnvironment> : RecordedTestBas
 
         protected ManagementRecordedTestBase(bool isAsync, RecordedTestMode? mode = default) : base(isAsync, mode)
         {
+            AdditionalInterceptors = new[] { new ManagementInterceptor(this) };
+
             SessionEnvironment = new TEnvironment();
             SessionEnvironment.Mode = Mode;
             Initialize();
@@ -224,8 +225,5 @@ public void OneTimeCleanupResourceGroups()
             if (!(GlobalClient is null))
                 throw new InvalidOperationException("StopSessionRecording was never called please make sure you call that at the end of your OneTimeSetup");
         }
-
-        protected override object InstrumentOperation(Type operationType, object operation)
-            => InstrumentOperationInternal(operationType, operation, new ManagementInterceptor(this));
     }
 }

sdk/communication/Azure.Communication.NetworkTraversal/tests/CommunicationRelayClient/CommunicationRelayClientLiveTests.cs

@@ -18,6 +18,7 @@ namespace Azure.Communication.NetworkTraversal.Tests
     /// These tests have a dependency on live Azure services and may incur costs for the associated
     /// Azure subscription.
     /// </remarks>
+    [Ignore("https://github.com/Azure/azure-sdk-for-net/issues/27522")]
     public class CommunicationRelayClientLiveTests : CommunicationRelayClientLiveTestBase
     {
         /// <summary>

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

@@ -4,6 +4,7 @@
 using System;
 using System.Collections.Generic;
 using System.Diagnostics;
+using System.Linq;
 using System.Reflection;
 using Castle.DynamicProxy;
 using NUnit.Framework;
@@ -36,6 +37,8 @@ public ClientTestBase(bool isAsync)
             IsAsync = isAsync;
         }
 
+        protected IReadOnlyCollection<IInterceptor> AdditionalInterceptors { get; set; }
+
         protected virtual DateTime TestStartTime => TestExecutionContext.CurrentContext.StartTime;
 
         [TearDown]
@@ -140,9 +143,20 @@ protected internal virtual object InstrumentClient(Type clientType, object clien
 
         protected internal virtual object InstrumentOperation(Type operationType, object operation)
         {
-            return operation;
+            var interceptors = AdditionalInterceptors ?? Array.Empty<IInterceptor>();
+
+            // The assumption is that any recorded or live tests deriving from RecordedTestBase, and that any unit tests deriving directly from ClientTestBase are equivalent to playback.
+            var interceptorArray = interceptors.Concat(new IInterceptor[] { new GetOriginalInterceptor(operation), new OperationInterceptor(RecordedTestMode.Playback) }).ToArray();
+            return ProxyGenerator.CreateClassProxyWithTarget(
+                operationType,
+                new[] { typeof(IInstrumented) },
+                operation,
+                interceptorArray);
         }
 
+        protected internal T InstrumentOperation<T>(T operation) where T : Operation =>
+            (T)InstrumentOperation(typeof(T), operation);
+
         protected T GetOriginal<T>(T instrumented)
         {
             if (instrumented == null) throw new ArgumentNullException(nameof(instrumented));

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

@@ -4,7 +4,6 @@
 using System;
 using System.Linq;
 using System.Reflection;
-using System.Runtime.ExceptionServices;
 using System.Threading;
 using Castle.DynamicProxy;
 
@@ -18,16 +17,16 @@ internal class OperationInterceptor : IInterceptor
 
         internal static readonly string PollerWaitForCompletionAsyncName = nameof(OperationPoller.WaitForCompletionAsync);
 
-        private readonly bool _noWait;
+        private readonly RecordedTestMode _mode;
 
-        public OperationInterceptor(bool noWait)
+        public OperationInterceptor(RecordedTestMode mode)
         {
-            _noWait = noWait;
+            _mode = mode;
         }
 
         public void Intercept(IInvocation invocation)
         {
-            if (_noWait)
+            if (_mode == RecordedTestMode.Playback)
             {
                 if (invocation.Method.Name == WaitForCompletionMethodName)
                 {

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

@@ -344,17 +344,10 @@ protected internal override object InstrumentClient(Type clientType, object clie
             return base.InstrumentClient(clientType, client, preInterceptors);
         }
 
-        protected internal T InstrumentOperation<T>(T operation) where T: Operation
-        {
-            return (T) InstrumentOperation(typeof(T), operation);
-        }
-
         protected internal override object InstrumentOperation(Type operationType, object operation)
-            => InstrumentOperationInternal(operationType, operation);
-
-        protected object InstrumentOperationInternal(Type operationType, object operation, params IInterceptor[] interceptors)
         {
-            var interceptorArray = interceptors.Concat(new IInterceptor[] { new GetOriginalInterceptor(operation), new OperationInterceptor(Mode == RecordedTestMode.Playback) }).ToArray();
+            var interceptors = AdditionalInterceptors ?? Array.Empty<IInterceptor>();
+            var interceptorArray = interceptors.Concat(new IInterceptor[] { new GetOriginalInterceptor(operation), new OperationInterceptor(Mode) }).ToArray();
             return ProxyGenerator.CreateClassProxyWithTarget(
                 operationType,
                 new[] { typeof(IInstrumented) },

sdk/keyvault/Azure.Security.KeyVault.Certificates/tests/CertificateClientLiveTests.cs

@@ -356,10 +356,11 @@ public async Task VerifyGetCertificateCompletedSubsequently()
 
             // Pretend a separate process was started subsequently and we need to get the operation again.
             CertificateOperation operation = new CertificateOperation(Client, certName);
+            operation = InstrumentOperation(operation);
 
             // Need to call the real async wait method or the sync version of this test fails because it's using the instrumented Client directly.
             using CancellationTokenSource cts = new CancellationTokenSource(DefaultCertificateOperationTimeout);
-            await operation.WaitForCompletionAsync(PollingInterval, cts.Token);
+            await operation.WaitForCompletionAsync(cts.Token);
 
             Assert.IsTrue(operation.HasCompleted);
             Assert.IsTrue(operation.HasValue);

sdk/keyvault/Azure.Security.KeyVault.Certificates/tests/CertificateOperationTests.cs

@@ -292,6 +292,11 @@ private CertificateClient CreateClient(HttpPipelineTransport transport)
         {
             CertificateClientOptions options = new CertificateClientOptions
             {
+                Retry =
+                {
+                    Delay = TimeSpan.FromMilliseconds(10),
+                    Mode = RetryMode.Fixed,
+                },
                 Transport = transport,
             };
 
@@ -303,10 +308,8 @@ private CertificateClient CreateClient(HttpPipelineTransport transport)
                     ));
         }
 
-        private async ValueTask<KeyVaultCertificateWithPolicy> WaitForOperationAsync(CertificateOperation operation)
-        {
-            return await operation.WaitForCompletionAsync(TimeSpan.Zero, default);
-        }
+        private async ValueTask<KeyVaultCertificateWithPolicy> WaitForOperationAsync(CertificateOperation operation) =>
+            await operation.WaitForCompletionAsync();
 
         public class MockCredential : TokenCredential
         {