diff --git a/sdk/core/Azure.Core.TestFramework/src/Instrumentation/InstrumentResultInterceptor.cs b/sdk/core/Azure.Core.TestFramework/src/Instrumentation/InstrumentResultInterceptor.cs index 90d317a849e3..bc089b96dc4d 100644 --- a/sdk/core/Azure.Core.TestFramework/src/Instrumentation/InstrumentResultInterceptor.cs +++ b/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 InstrumentOperationInterceptor(IInvocation invocation, Func> innerTask) { return (T) _testBase.InstrumentOperation(typeof(T), await innerTask()); diff --git a/sdk/core/Azure.Core.TestFramework/src/Shared/TestClient.cs b/sdk/core/Azure.Core.TestFramework/src/Shared/TestClient.cs index c20de1d9de0d..9f7ae298776e 100644 --- a/sdk/core/Azure.Core.TestFramework/src/Shared/TestClient.cs +++ b/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 MethodBAsync() return nameof(MethodAAsync); } } + +#pragma warning disable SA1402 + internal class AnotherType +#pragma warning restore SA1402 + { + public virtual HttpPipeline Pipeline { get; } + + public virtual Task 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; + } + } } diff --git a/sdk/core/Azure.Core.TestFramework/tests/ClientTestBaseTests.cs b/sdk/core/Azure.Core.TestFramework/tests/ClientTestBaseTests.cs index fc6de3d665f1..5d7d9788d217 100644 --- a/sdk/core/Azure.Core.TestFramework/tests/ClientTestBaseTests.cs +++ b/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() {