Update tests and framework to instrument all clients recursively

common/ManagementTestShared/Redesign/ManagementRecordedTestBase.cs

@@ -63,7 +63,7 @@ protected AzureResourceManagerClient GetArmClient()
         }
 
         [SetUp]
-        protected void Setup()
+        protected void CreateCleanupClient()
         {
             _cleanupClient ??= GetCleanupClient();
         }
@@ -75,7 +75,14 @@ protected void CleanupResourceGroups()
             {
                 Parallel.ForEach(CleanupPolicy.ResourceGroupsCreated, resourceGroup =>
                 {
-                    _cleanupClient.GetResourceGroupOperations(TestEnvironment.SubscriptionId, resourceGroup).StartDelete();
+                    try
+                    {
+                        _cleanupClient.GetResourceGroupOperations(TestEnvironment.SubscriptionId, resourceGroup).StartDelete();
+                    }
+                    catch (RequestFailedException e) when (e.Status == 404)
+                    {
+                        //we assume the test case cleaned up it up if it no longer exists.
+                    }
                 });
             }
         }
@@ -136,7 +143,7 @@ private bool HasOneTimeSetup()
             var methods = GetType().GetMethods().Where(m => types.Contains(m.DeclaringType));
             foreach (var method in methods)
             {
-                foreach(var attr in method.GetCustomAttributes(false))
+                foreach (var attr in method.GetCustomAttributes(false))
                 {
                     if (attr is OneTimeSetUpAttribute)
                         return true;

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

@@ -0,0 +1,34 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+using System.Collections.Generic;
+using System.Threading.Tasks;
+using Castle.DynamicProxy;
+
+namespace Azure.Core.TestFramework
+{
+    public class AsyncPageableInterceptor<T> : IAsyncEnumerator<T>
+        where T : class
+    {
+        private ClientTestBase _testBase;
+        private IAsyncEnumerator<T> _inner;
+
+        public AsyncPageableInterceptor(ClientTestBase testBase, IAsyncEnumerator<T> inner)
+        {
+            _testBase = testBase;
+            _inner = inner;
+        }
+
+        public T Current => _testBase.InstrumentClient(typeof(T), _inner.Current, new IInterceptor[] { new ManagementInterceptor(_testBase) }) as T;
+
+        public ValueTask<bool> MoveNextAsync()
+        {
+            return _inner.MoveNextAsync();
+        }
+
+        public ValueTask DisposeAsync()
+        {
+            return _inner.DisposeAsync();
+        }
+    }
+}

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

@@ -33,7 +33,9 @@ public void Intercept(IInvocation invocation)
                 invocation.Proceed();
                 invocation.ReturnValue = Activator.CreateInstance(typeof(DiagnosticScopeValidatingAsyncEnumerable<>).MakeGenericType(genericType.GenericTypeArguments[0]), invocation.ReturnValue, expectedName, methodName, strict);
             }
-            else if (methodName.EndsWith("Async") && !invocation.Method.ReturnType.Name.Contains("IAsyncEnumerable"))
+            else if (methodName.EndsWith("Async") &&
+                !invocation.Method.ReturnType.Name.Contains("IAsyncEnumerable") &&
+                !invocation.Method.Name.Contains("WaitForCompletionAsync"))
             {
                 Type genericArgument = typeof(object);
                 Type awaitableType = invocation.Method.ReturnType;

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

@@ -26,14 +26,24 @@ public void Intercept(IInvocation invocation)
             }
 
             var type = result.GetType();
-            // We don't want to instrument generated rest clients.
-            if ((type.Name.EndsWith("Client") && !type.Name.EndsWith("RestClient")) ||
+
+            if (
+                // We don't want to instrument generated rest clients.
+                type.Name.EndsWith("Client") && !type.Name.EndsWith("RestClient"))
+            {
+                invocation.ReturnValue = _testBase.InstrumentClient(type, result, Array.Empty<IInterceptor>());
+                return;
+            }
+
+            if (
                 // Generated ARM clients will have a property containing the sub-client that ends with Operations.
                 (invocation.Method.Name.StartsWith("get_") && type.Name.EndsWith("Operations")) ||
-                // Instrument the subscription client that hangs off of the new AzureResouceManagementClient
-                (type.Name.EndsWith("DefaultSubscription")))
+                // Instrument the container construction methods inside Operations objects
+                (invocation.Method.Name.StartsWith("Get") && type.Name.EndsWith("Container")) ||
+                // Instrument the operations construction methods inside Operations objects
+                (invocation.Method.Name.StartsWith("Get") && type.Name.EndsWith("Operations")))
             {
-                invocation.ReturnValue = _testBase.InstrumentClient(type, result, Array.Empty<IInterceptor>());
+                invocation.ReturnValue = _testBase.InstrumentClient(type, result, new IInterceptor[] { new ManagementInterceptor(_testBase) });
             }
         }
     }

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

@@ -0,0 +1,63 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+using System;
+using System.Reflection;
+using System.Threading.Tasks;
+using Castle.DynamicProxy;
+
+namespace Azure.Core.TestFramework
+{
+    internal class ManagementInterceptor : IInterceptor
+    {
+        private readonly ClientTestBase _testBase;
+        private static readonly ProxyGenerator s_proxyGenerator = new ProxyGenerator();
+
+        public ManagementInterceptor(ClientTestBase testBase)
+        {
+            _testBase = testBase;
+        }
+
+        public void Intercept(IInvocation invocation)
+        {
+            invocation.Proceed();
+
+            var result = invocation.ReturnValue;
+            if (result == null)
+            {
+                return;
+            }
+
+            var type = result.GetType();
+
+            if (type.Name.StartsWith("Task"))
+            {
+                var taskResultType = type.GetGenericArguments()[0];
+                if (taskResultType.Name.StartsWith("ArmResponse") || taskResultType.Name.StartsWith("ArmOperation"))
+                {
+                    var taskResult = result.GetType().GetProperty("Result").GetValue(result);
+                    var instrumentedResult = _testBase.InstrumentClient(taskResultType, taskResult, new IInterceptor[] { new ManagementInterceptor(_testBase) });
+                    var method = typeof(Task).GetMethod("FromResult", BindingFlags.Public | BindingFlags.Static);
+                    var genericType = taskResultType.Name.StartsWith("Ph") ? taskResultType.BaseType : taskResultType; //TODO: remove after 5279 and 5284
+                    var genericMethod = method.MakeGenericMethod(genericType);
+                    invocation.ReturnValue = genericMethod.Invoke(null, new object[] { instrumentedResult });
+                }
+            }
+            else if (invocation.Method.Name.EndsWith("Value") && type.BaseType.Name.EndsWith("Operations"))
+            {
+                invocation.ReturnValue = _testBase.InstrumentClient(type, result, new IInterceptor[] { new ManagementInterceptor(_testBase) });
+            }
+            else if (type.BaseType.Name.StartsWith("AsyncPageable"))
+            {
+                invocation.ReturnValue = s_proxyGenerator.CreateClassProxyWithTarget(type, result, new IInterceptor[] { new ManagementInterceptor(_testBase) });
+            }
+            else if (invocation.Method.Name.StartsWith("Get") && invocation.Method.Name.EndsWith("Enumerator"))
+            {
+                var wrapperType = typeof(AsyncPageableInterceptor<>);
+                var genericType = wrapperType.MakeGenericType(type.GetGenericArguments()[0]);
+                var ctor = genericType.GetConstructor(new Type[] { typeof(ClientTestBase), result.GetType() });
+                invocation.ReturnValue = ctor.Invoke(new object[] { _testBase, result });
+            }
+        }
+    }
+}

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

@@ -247,10 +247,14 @@ MethodInfo GetMethodSlow()
 
         private static bool IsInternal(MethodBase method) => method.IsAssembly || method.IsFamilyAndAssembly && !method.IsFamilyOrAssembly;
 
-        private class SyncPageableWrapper<T> : AsyncPageable<T>
+        public class SyncPageableWrapper<T> : AsyncPageable<T>
         {
             private readonly Pageable<T> _enumerable;
 
+            protected SyncPageableWrapper()
+            {
+            }
+
             public SyncPageableWrapper(Pageable<T> enumerable)
             {
                 _enumerable = enumerable;

sdk/core/Azure.Core/Azure.Core.sln

@@ -21,11 +21,15 @@ Project("{9A19103F-16F7-4668-BE54-9A1E7A4F7556}") = "Microsoft.Azure.Core.Spatia
 EndProject
 Project("{9A19103F-16F7-4668-BE54-9A1E7A4F7556}") = "Microsoft.Azure.Core.Spatial.NewtonsoftJson.Tests", "..\Microsoft.Azure.Core.Spatial.NewtonsoftJson\tests\Microsoft.Azure.Core.Spatial.NewtonsoftJson.Tests.csproj", "{622772C8-A2CB-4F8B-82EB-2A46BCC21DAE}"
 EndProject
-Project("{FAE04EC0-301F-11D3-BF4B-00C04F79EFBC}") = "Microsoft.Azure.Core.Spatial", "..\Microsoft.Azure.Core.Spatial\src\Microsoft.Azure.Core.Spatial.csproj", "{25A7E209-D5D2-41F1-AFE9-E860D6294EF5}"
+Project("{9A19103F-16F7-4668-BE54-9A1E7A4F7556}") = "Microsoft.Azure.Core.Spatial", "..\Microsoft.Azure.Core.Spatial\src\Microsoft.Azure.Core.Spatial.csproj", "{25A7E209-D5D2-41F1-AFE9-E860D6294EF5}"
 EndProject
-Project("{FAE04EC0-301F-11D3-BF4B-00C04F79EFBC}") = "Microsoft.Azure.Core.Spatial.Tests", "..\Microsoft.Azure.Core.Spatial\tests\Microsoft.Azure.Core.Spatial.Tests.csproj", "{36FF59A9-D6C9-4226-A0FE-47C879F3EB94}"
+Project("{9A19103F-16F7-4668-BE54-9A1E7A4F7556}") = "Microsoft.Azure.Core.Spatial.Tests", "..\Microsoft.Azure.Core.Spatial\tests\Microsoft.Azure.Core.Spatial.Tests.csproj", "{36FF59A9-D6C9-4226-A0FE-47C879F3EB94}"
 EndProject
-Project("{FAE04EC0-301F-11D3-BF4B-00C04F79EFBC}") = "Azure.Core.Perf", "perf\Azure.Core.Perf.csproj", "{B6D7909F-4DE6-4895-BF39-EF892BA64BA3}"
+Project("{9A19103F-16F7-4668-BE54-9A1E7A4F7556}") = "Azure.Core.Perf", "perf\Azure.Core.Perf.csproj", "{B6D7909F-4DE6-4895-BF39-EF892BA64BA3}"
+EndProject
+Project("{9A19103F-16F7-4668-BE54-9A1E7A4F7556}") = "Azure.Test.Perf", "..\..\..\common\Perf\Azure.Test.Perf\Azure.Test.Perf.csproj", "{96E9F605-9C38-4D77-96F4-679EF8B9390F}"
+EndProject
+Project("{9A19103F-16F7-4668-BE54-9A1E7A4F7556}") = "Azure.ResourceManager.Core", "..\..\resourcemanager\Azure.ResourceManager.Core\src\Azure.ResourceManager.Core.csproj", "{8E60A748-3973-471A-B103-EC9406BB3313}"
 EndProject
 Global
 	GlobalSection(SolutionConfigurationPlatforms) = preSolution
@@ -81,6 +85,14 @@ Global
 		{B6D7909F-4DE6-4895-BF39-EF892BA64BA3}.Debug|Any CPU.Build.0 = Debug|Any CPU
 		{B6D7909F-4DE6-4895-BF39-EF892BA64BA3}.Release|Any CPU.ActiveCfg = Release|Any CPU
 		{B6D7909F-4DE6-4895-BF39-EF892BA64BA3}.Release|Any CPU.Build.0 = Release|Any CPU
+		{96E9F605-9C38-4D77-96F4-679EF8B9390F}.Debug|Any CPU.ActiveCfg = Debug|Any CPU
+		{96E9F605-9C38-4D77-96F4-679EF8B9390F}.Debug|Any CPU.Build.0 = Debug|Any CPU
+		{96E9F605-9C38-4D77-96F4-679EF8B9390F}.Release|Any CPU.ActiveCfg = Release|Any CPU
+		{96E9F605-9C38-4D77-96F4-679EF8B9390F}.Release|Any CPU.Build.0 = Release|Any CPU
+		{8E60A748-3973-471A-B103-EC9406BB3313}.Debug|Any CPU.ActiveCfg = Debug|Any CPU
+		{8E60A748-3973-471A-B103-EC9406BB3313}.Debug|Any CPU.Build.0 = Debug|Any CPU
+		{8E60A748-3973-471A-B103-EC9406BB3313}.Release|Any CPU.ActiveCfg = Release|Any CPU
+		{8E60A748-3973-471A-B103-EC9406BB3313}.Release|Any CPU.Build.0 = Release|Any CPU
 	EndGlobalSection
 	GlobalSection(SolutionProperties) = preSolution
 		HideSolutionNode = FALSE

sdk/core/Azure.Core/tests/Azure.Core.Tests.csproj

@@ -16,6 +16,7 @@
   </ItemGroup>
   <ItemGroup>
     <ProjectReference Include="$(AzureCoreTestFramework)" />
+    <ProjectReference Include="..\..\..\resourcemanager\Azure.ResourceManager.Core\src\Azure.ResourceManager.Core.csproj" />
     <ProjectReference Include="..\src\Azure.Core.csproj" />
     <ProjectReference Include="..\..\Microsoft.Azure.Core.NewtonsoftJson\src\Microsoft.Azure.Core.NewtonsoftJson.csproj" />
     <PackageReference Include="Azure.Security.KeyVault.Secrets" />
@@ -39,5 +40,6 @@
     <Compile Include="..\src\Shared\ValueStopwatch.cs" LinkBase="Shared" />
     <Compile Include="..\src\Shared\LightweightPkcs8Decoder.cs" LinkBase="Shared" />
     <Compile Include="..\src\Shared\PemReader.cs" LinkBase="Shared" />
+    <Compile Include="..\..\..\..\common\ManagementTestShared\Redesign\*.cs" LinkBase="Shared\Mgmt" />
   </ItemGroup>
 </Project>

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

@@ -160,7 +160,7 @@ public async Task SubClientPropertyCallsAreAutoInstrumented()
         {
             TestClient client = InstrumentClient(new TestClient());
 
-            Operations subClient = client.SubProperty;
+            TestClientOperations subClient = client.SubProperty;
             var result = await subClient.MethodAsync(123);
 
             Assert.AreEqual(IsAsync ? "Async 123 False" : "Sync 123 False", result);
@@ -191,145 +191,5 @@ public async Task TasksValidateOwnScopes()
             });
             await Task.WhenAll(t1, t2);
         }
-
-        public class TestClient
-        {
-            private readonly ClientDiagnostics _diagnostics;
-
-            public TestClient() : this(null)
-            {
-            }
-
-            public TestClient(TestClientOptions options)
-            {
-                options ??= new TestClientOptions();
-                _diagnostics = new ClientDiagnostics(options);
-            }
-
-            public virtual Task<string> MethodAsync(int i, CancellationToken cancellationToken = default)
-            {
-                return Task.FromResult("Async " + i + " " + cancellationToken.CanBeCanceled);
-            }
-
-            public virtual Task<string> MethodGenericAsync<T>(T i, CancellationToken cancellationToken = default)
-            {
-                return Task.FromResult($"Async {i} {cancellationToken.CanBeCanceled}");
-            }
-
-            public virtual string MethodGeneric<T>(T i, CancellationToken cancellationToken = default)
-            {
-                return $"Sync {i} {cancellationToken.CanBeCanceled}";
-            }
-
-            public virtual Task<string> NoAlternativeAsync(int i, CancellationToken cancellationToken = default)
-            {
-                return Task.FromResult("I don't have sync alternative");
-            }
-
-            public virtual string Method(int i, CancellationToken cancellationToken = default)
-            {
-                return "Sync " + i + " " + cancellationToken.CanBeCanceled;
-            }
-
-            public virtual string Method2()
-            {
-                return "Hello";
-            }
-
-            // These four follow the new pattern for custom users schemas
-            public virtual Task<Response<T>> GetDataAsync<T>() =>
-                Task.FromResult(Response.FromValue(default(T), new MockResponse(200, "async - static")));
-            public virtual Response<T> GetData<T>(T arg) =>
-                Response.FromValue(default(T), new MockResponse(200, $"sync - static {arg}"));
-            public virtual Task<Response<T>> GetDataAsync<T>(T arg) =>
-                Task.FromResult(Response.FromValue(default(T), new MockResponse(200, $"async - static {arg}")));
-            public virtual Response<T> GetData<T>() =>
-                Response.FromValue(default(T), new MockResponse(200, "sync - static"));
-            public virtual Task<Response<object>> GetDataAsync() =>
-                Task.FromResult(Response.FromValue((object)null, new MockResponse(200, "async - dynamic")));
-            public virtual Response<object> GetData() =>
-                Response.FromValue((object)null, new MockResponse(200, "sync - dynamic"));
-
-            // These four follow the new pattern for custom users schemas and
-            // throw exceptions
-            public virtual Task<Response<T>> GetFailureAsync<T>() =>
-                throw new InvalidOperationException("async - static");
-            public virtual Response<T> GetFailure<T>() =>
-                throw new InvalidOperationException("sync - static");
-            public virtual Task<Response<object>> GetFailureAsync() =>
-                throw new InvalidOperationException("async - dynamic");
-            public virtual Response<object> GetFailure() =>
-                throw new InvalidOperationException("sync - dynamic");
-
-            public virtual TestClient GetAnotherTestClient()
-            {
-                return new TestClient();
-            }
-            public virtual Operations SubProperty => new Operations();
-
-            public virtual string MethodA()
-            {
-                using DiagnosticScope scope = _diagnostics.CreateScope($"{nameof(TestClient)}.{nameof(MethodA)}");
-                scope.Start();
-
-                return nameof(MethodA);
-            }
-
-            public virtual async Task<string> MethodAAsync()
-            {
-                using DiagnosticScope scope = _diagnostics.CreateScope($"{nameof(TestClient)}.{nameof(MethodA)}");
-                scope.Start();
-
-                await Task.Yield();
-                return nameof(MethodAAsync);
-            }
-
-            public virtual string MethodB()
-            {
-                using DiagnosticScope scope = _diagnostics.CreateScope($"{nameof(TestClient)}.{nameof(MethodB)}");
-                scope.Start();
-
-                return nameof(MethodB);
-            }
-
-            public virtual async Task<string> MethodBAsync()
-            {
-                using DiagnosticScope scope = _diagnostics.CreateScope($"{nameof(TestClient)}.{nameof(MethodB)}");
-                scope.Start();
-
-                await Task.Yield();
-                return nameof(MethodAAsync);
-            }
-        }
-
-        public class TestClientOptions : ClientOptions
-        {
-        }
-
-        public class Operations
-        {
-            public virtual Task<string> MethodAsync(int i, CancellationToken cancellationToken = default)
-            {
-                return Task.FromResult("Async " + i + " " + cancellationToken.CanBeCanceled);
-            }
-
-            public virtual string Method(int i, CancellationToken cancellationToken = default)
-            {
-                return "Sync " + i + " " + cancellationToken.CanBeCanceled;
-            }
-        }
-
-        public class InvalidTestClient
-        {
-            public Task<string> MethodAsync(int i)
-            {
-                return Task.FromResult("Async " + i);
-            }
-
-            public virtual string Method(int i)
-            {
-                return "Sync " + i;
-            }
-        }
     }
 }