Update analyzer to support [TestingAndPrivateOnly] attribute

Author: andrewlock

tracer/src/Datadog.Trace.Tools.Analyzers/TestingOnlyAnalyzer/TestingOnlyAnalyzer.cs

@@ -19,9 +19,10 @@ namespace Datadog.Trace.Tools.Analyzers.TestingOnlyAnalyzer
     /// <summary>
     /// DD002: Incorrect usage of internal API
     ///
-    /// Finds internal usages of APIs specifically marked with the [TestingOnly] flag.
-    /// These methods should not be called directly by our library code, only from test code.
-    /// The analyzer enforces that requirement
+    /// Finds internal usages of APIs specifically marked with the [TestingOnly] or [TestingAndPrivateOnly] flag.
+    /// - [TestingOnly]: Methods should not be called directly by our library code, only from test code.
+    /// - [TestingAndPrivateOnly]: Methods can be called from within the same type, or from test code, but not from other types.
+    /// The analyzer enforces these requirements.
     ///
     /// </summary>
     [DiagnosticAnalyzer(LanguageNames.CSharp)]
@@ -33,9 +34,7 @@ public sealed class TestingOnlyAnalyzer : DiagnosticAnalyzer
         public const string DiagnosticId = "DD0002";
 
         private const string TestingOnlyAttribute = nameof(TestingOnlyAttribute);
-
-        private static readonly ImmutableArray<string> TestingOnlyAttributeNames
-            = ImmutableArray.Create(TestingOnlyAttribute);
+        private const string TestingAndPrivateOnlyAttribute = nameof(TestingAndPrivateOnlyAttribute);
 
 #pragma warning disable RS2008 // Enable analyzer release tracking for the analyzer project
         private static readonly DiagnosticDescriptor Rule = new(
@@ -71,11 +70,12 @@ private static void AnalyzeOperationBlock(
             ConcurrentDictionary<ISymbol, InternalApiStatus> internalApiMembers)
         {
             var internalApiOperations = PooledConcurrentDictionary<KeyValuePair<IOperation, ISymbol>, bool>.GetInstance();
+            var callingType = context.OwningSymbol?.ContainingType;
 
             context.RegisterOperationAction(
                 ctx =>
                 {
-                    Helpers.AnalyzeOperation(ctx.Operation, internalApiOperations, internalApiMembers);
+                    Helpers.AnalyzeOperation(ctx.Operation, internalApiOperations, internalApiMembers, callingType);
                 },
                 OperationKind.MethodReference,
                 OperationKind.EventReference,
@@ -110,7 +110,8 @@ private static class Helpers
             internal static void AnalyzeOperation(
                 IOperation operation,
                 PooledConcurrentDictionary<KeyValuePair<IOperation, ISymbol>, bool> internalApiOperations,
-                ConcurrentDictionary<ISymbol, InternalApiStatus> internalApiMembers)
+                ConcurrentDictionary<ISymbol, InternalApiStatus> internalApiMembers,
+                INamedTypeSymbol? callingType)
             {
                 var symbol = GetOperationSymbol(operation);
 
@@ -185,8 +186,21 @@ void CheckTypeArgumentsCore(ImmutableArray<ITypeSymbol> typeArguments, PooledHas
 
                 void CheckOperationAttributes(ISymbol symbol, bool checkParents)
                 {
-                    if (TryGetOrCreatePublicApiAttributes(symbol, checkParents, internalApiMembers, out _))
+                    if (TryGetOrCreatePublicApiAttributes(symbol, checkParents, internalApiMembers, out var attributes))
                     {
+                        // If the marked member has [TestingAndPrivateOnly], allow calls from within the same type
+                        if (attributes.IsTestingAndPrivateOnly)
+                        {
+                            var targetType = symbol.ContainingType;
+
+                            // Allow calls from within the same type
+                            if (callingType != null && targetType != null &&
+                                SymbolEqualityComparer.Default.Equals(callingType, targetType))
+                            {
+                                return;
+                            }
+                        }
+
                         internalApiOperations.TryAdd(new KeyValuePair<IOperation, ISymbol>(operation, symbol), true);
                     }
                 }
@@ -243,7 +257,8 @@ private static InternalApiStatus CopyAttributes(InternalApiStatus copyAttributes
                 new()
                 {
                     IsAssemblyAttribute = copyAttributes.IsAssemblyAttribute,
-                    IsPublicApi = copyAttributes.IsPublicApi
+                    IsTestingApi = copyAttributes.IsTestingApi,
+                    IsTestingAndPrivateOnly = copyAttributes.IsTestingAndPrivateOnly
                 };
 
             private static bool IsWithinConditionalOperation(IFieldReferenceOperation pOperation) =>
@@ -287,17 +302,25 @@ private static bool TryGetOrCreatePublicApiAttributes(
                     attributes = publicApiMembers.GetOrAdd(symbol, attributes);
                 }
 
-                return attributes.IsPublicApi;
+                return attributes.IsTestingApi;
 
                 static void MergePlatformAttributes(
                     ImmutableArray<AttributeData> immediateAttributes,
                     ref InternalApiStatus parentAttributes)
                 {
                     foreach (AttributeData attribute in immediateAttributes)
                     {
-                        if (TestingOnlyAttributeNames.Contains(attribute.AttributeClass!.Name))
+                        var attributeName = attribute.AttributeClass!.Name;
+                        if (attributeName == TestingAndPrivateOnlyAttribute)
                         {
-                            parentAttributes.IsPublicApi = true;
+                            parentAttributes.IsTestingApi = true;
+                            parentAttributes.IsTestingAndPrivateOnly = true;
+                            return;
+                        }
+                        else if (attributeName == TestingOnlyAttribute)
+                        {
+                            parentAttributes.IsTestingApi = true;
+                            parentAttributes.IsTestingAndPrivateOnly = false;
                             return;
                         }
                     }
@@ -309,7 +332,9 @@ private sealed class InternalApiStatus
         {
             public bool IsAssemblyAttribute { get; set; }
 
-            public bool IsPublicApi { get; set; }
+            public bool IsTestingApi { get; set; }
+
+            public bool IsTestingAndPrivateOnly { get; set; }
         }
     }
 }

tracer/test/Datadog.Trace.Tools.Analyzers.Tests/TestingOnlyAnalyzerTests.cs

@@ -23,6 +23,8 @@ public class TestingOnlyAnalyzerTests
 
         public static string[] GetTestingOnlyAttributes { get; } = { "TestingOnly", "TestingOnlyAttribute" };
 
+        public static string[] GetTestingAndPrivateOnlyAttributes { get; } = { "TestingAndPrivateOnly", "TestingAndPrivateOnlyAttribute" };
+
         public static string[] NonTestingOnlyAccesses { get; } =
         {
             "var x = _nonPublicField;",
@@ -65,6 +67,26 @@ public class TestingOnlyAnalyzerTests
             """,
         };
 
+        public static string[] TestingAndPrivateOnlySameTypeAccesses { get; } =
+        {
+            // Access members from the same class - should NOT be flagged
+            "var x = _publicField;",
+            "var x = PublicProperty;",
+            "PublicProperty = string.Empty;",
+            "var x = PublicMethod();",
+            "var x = new TestClass();",
+        };
+
+        public static string[] TestingAndPrivateOnlyDifferentTypeAccesses { get; } =
+        {
+            // Access members from a different class - should be flagged
+            "var x = {|#0:OtherClass._testingAndPrivateField|};",
+            "var x = {|#0:OtherClass.TestingAndPrivateProperty|};",
+            "{|#0:OtherClass.TestingAndPrivateProperty|} = string.Empty;",
+            "var x = {|#0:OtherClass.TestingAndPrivateMethod()|};",
+            "var x = {|#0:new OtherClass()|};",
+        };
+
         public static IEnumerable<object[]> NonPublicCombination { get; } =
             from attrs in GetTestingOnlyAttributes
             from includeNamespace in new[] { true, false }
@@ -86,6 +108,20 @@ from api in ShouldThrowButDoesnt
             from conditional in new[] { true, false }
             select new object[] { attrs, includeNamespace, api, conditional };
 
+        public static IEnumerable<object[]> TestingAndPrivateOnlySameTypeCombination { get; } =
+            from attrs in GetTestingAndPrivateOnlyAttributes
+            from includeNamespace in new[] { true, false }
+            from api in TestingAndPrivateOnlySameTypeAccesses
+            from conditional in new[] { true, false }
+            select new object[] { attrs, includeNamespace, api, conditional };
+
+        public static IEnumerable<object[]> TestingAndPrivateOnlyDifferentTypeCombination { get; } =
+            from attrs in GetTestingAndPrivateOnlyAttributes
+            from includeNamespace in new[] { true, false }
+            from api in TestingAndPrivateOnlyDifferentTypeAccesses
+            from conditional in new[] { true, false }
+            select new object[] { attrs, includeNamespace, api, conditional };
+
         [Fact]
         public async Task EmptySourceShouldNotHaveDiagnostics()
         {
@@ -133,6 +169,27 @@ public async Task NotSupported(string publicAttribute, bool includeNamespace, st
             await ideallyShouldThrow.Should().ThrowAsync<InvalidOperationException>();
         }
 
+        [Theory]
+        [MemberData(nameof(TestingAndPrivateOnlySameTypeCombination))]
+        public async Task ShouldNotFlagTestingAndPrivateOnlyWhenCalledFromSameType(string publicAttribute, bool includeNamespace, string testFragment, bool attributeIsConditional)
+        {
+            var code = GetSampleCode(publicAttribute, includeNamespace, testFragment, attributeIsConditional);
+
+            // No diagnostics expected - calls from within the same type should be allowed
+            await Verifier.VerifyAnalyzerAsync(code);
+        }
+
+        [Theory]
+        [MemberData(nameof(TestingAndPrivateOnlyDifferentTypeCombination))]
+        public async Task ShouldFlagTestingAndPrivateOnlyWhenCalledFromDifferentType(string publicAttribute, bool includeNamespace, string testFragment, bool attributeIsConditional)
+        {
+            var code = GetSampleCode(publicAttribute, includeNamespace, testFragment, attributeIsConditional);
+
+            var expected = new DiagnosticResult(DiagnosticId, DiagnosticSeverity.Error)
+               .WithLocation(0);
+            await Verifier.VerifyAnalyzerAsync(code, expected);
+        }
+
         private static string GetSampleCode(string publicAttribute, bool includeNamespace, string testFragment, bool attributeIsConditional)
         {
             var attributePrefix = includeNamespace ? "ConsoleApplication1." : string.Empty;
@@ -147,7 +204,7 @@ private static string GetSampleCode(string publicAttribute, bool includeNamespac
                 using System.Diagnostics;
 
                 namespace ConsoleApplication1;
-                
+
                 {{GetAttributeDefinition(publicAttribute, attributeIsConditional)}}
 
                 class TestClass
@@ -161,7 +218,7 @@ internal TestClass(string arg1) { }
 
                     [{{attributePrefix}}{{publicAttribute}}]
                     public TestClass() { }
-                    
+
                     public string NonPublicProperty { get; set; }
 
                     [{{attributePrefix}}{{publicAttribute}}]
@@ -211,6 +268,22 @@ public PublicClass() { }
                     public static string PublicMethod() => "test";
                 }
 
+                // Class for testing different-type access with TestingAndPrivateOnly
+                class OtherClass
+                {
+                    [{{attributePrefix}}{{publicAttribute}}]
+                    public static string _testingAndPrivateField;
+
+                    [{{attributePrefix}}{{publicAttribute}}]
+                    public OtherClass() { }
+
+                    [{{attributePrefix}}{{publicAttribute}}]
+                    public static string TestingAndPrivateProperty { get; set; }
+
+                    [{{attributePrefix}}{{publicAttribute}}]
+                    public static string TestingAndPrivateMethod() => "test";
+                }
+
                 interface IPublic
                 {
                     string PublicMethod();