Add subclient property interceptor support (Azure#53111)

* Add subclient property interceptor support

* fix

* Add test

* fix test

* Handle both property and method

* comment

* refactor

* fix

Author: JoshLove-msft

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

@@ -31,10 +31,7 @@ public void Intercept(IInvocation invocation)
             var type = invocation.Method.ReturnType;
 
             // We don't want to instrument generated rest clients.
-            if ((type.Name.EndsWith("Client") && !type.Name.EndsWith("RestClient") && !type.Name.EndsWith("ExtensionClient")) ||
-                // Generated ARM clients will have a property containing the sub-client that ends with Operations.
-                //TODO: remove after all track2 .net mgmt libraries are updated to the new generation
-                (invocation.Method.Name.StartsWith("get_") && type.Name.EndsWith("Operations")))
+            if (IsInstrumentableClientType(invocation))
             {
                 if (IsNullResult(invocation))
                     return;
@@ -100,6 +97,38 @@ public void Intercept(IInvocation invocation)
             invocation.Proceed();
         }
 
+        private static bool IsInstrumentableClientType(IInvocation invocation)
+        {
+            var type = invocation.Method.ReturnType;
+
+            // Skip system types
+            if (type.Namespace?.StartsWith("System.") == true)
+            {
+                return false;
+            }
+
+            // We don't want to instrument generated rest clients or extension clients
+            if (type.Name.EndsWith("RestClient") || type.Name.EndsWith("ExtensionClient"))
+            {
+                return false;
+            }
+
+            if (type.Name.EndsWith("Client"))
+            {
+                return true;
+            }
+
+            //TODO: remove after all track2 .net mgmt libraries are updated to the new generation
+            if (invocation.Method.Name.StartsWith("get_") && type.Name.EndsWith("Operations"))
+            {
+                return true;
+            }
+
+            // Some libraries have subclients that do not end with Client. Instrument any type that has a Pipeline property.
+            // This is the most expensive check so we do it last.
+            return type.GetProperty("Pipeline") != null;
+        }
+
         internal async ValueTask<T> InstrumentOperationInterceptor<T>(IInvocation invocation, Func<ValueTask<T>> innerTask)
         {
             return (T) _testBase.InstrumentOperation(typeof(T), await innerTask());

sdk/core/Azure.Core.TestFramework/src/Shared/TestClient.cs

@@ -85,6 +85,10 @@ public virtual TestClient GetAnotherTestClient()
         }
         public virtual TestClientOperations SubProperty => new TestClientOperations();
 
+        public virtual AnotherType SubClientProperty => new AnotherType();
+
+        public virtual AnotherType GetAnotherType() => new AnotherType();
+
         public virtual string MethodA()
         {
             using DiagnosticScope scope = _diagnostics.CreateScope($"{nameof(TestClient)}.{nameof(MethodA)}");
@@ -119,4 +123,21 @@ public virtual async Task<string> MethodBAsync()
             return nameof(MethodAAsync);
         }
     }
+
+#pragma warning disable SA1402
+    internal class AnotherType
+#pragma warning restore SA1402
+    {
+        public virtual HttpPipeline Pipeline { get; }
+
+        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;
+        }
+    }
 }

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

@@ -156,6 +156,28 @@ public async Task GetClientCallsAreAutoInstrumented()
             Assert.AreEqual(IsAsync ? "Async 123 False" : "Sync 123 False", result);
         }
 
+        [Test]
+        public async Task SubClientPropertyWithoutClientSuffixIsAutoInstrumented()
+        {
+            TestClient client = InstrumentClient(new TestClient());
+
+            AnotherType subClient = client.SubClientProperty;
+            var result = await subClient.MethodAsync(123);
+
+            Assert.AreEqual(IsAsync ? "Async 123 False" : "Sync 123 False", result);
+        }
+
+        [Test]
+        public async Task SubClientWithoutClientSuffixIsAutoInstrumented()
+        {
+            TestClient client = InstrumentClient(new TestClient());
+
+            AnotherType subClient = client.GetAnotherType();
+            var result = await subClient.MethodAsync(123);
+
+            Assert.AreEqual(IsAsync ? "Async 123 False" : "Sync 123 False", result);
+        }
+
         [Test]
         public async Task SubClientPropertyCallsAreAutoInstrumented()
         {