Fix MSTest fixture method fixers removing comments from method bodies

src/Analyzers/MSTest.Analyzers.CodeFixes/Helpers/FixtureMethodFixer.cs

@@ -5,19 +5,24 @@
 
 using Microsoft.CodeAnalysis;
 using Microsoft.CodeAnalysis.CSharp;
+using Microsoft.CodeAnalysis.CSharp.Syntax;
 using Microsoft.CodeAnalysis.Editing;
 
 namespace MSTest.Analyzers.Helpers;
 
 internal static class FixtureMethodFixer
 {
-    private const SyntaxNode? VoidReturnTypeNode = null;
-
     public static async Task<Solution> FixSignatureAsync(Document document, SyntaxNode root, SyntaxNode node,
         bool isParameterLess, bool shouldBeStatic, CancellationToken cancellationToken)
     {
         cancellationToken.ThrowIfCancellationRequested();
 
+        // Cast to MethodDeclarationSyntax to preserve method body and trivia
+        if (node is not MethodDeclarationSyntax methodDeclaration)
+        {
+            return document.Project.Solution;
+        }
+
         SemanticModel? semanticModel = await document.GetRequiredSemanticModelAsync(cancellationToken).ConfigureAwait(false);
 
         var methodSymbol = (IMethodSymbol?)semanticModel.GetDeclaredSymbol(node, cancellationToken);
@@ -27,66 +32,111 @@
         }
 
         var wellKnownTypeProvider = WellKnownTypeProvider.GetOrCreate(semanticModel.Compilation);
-        var syntaxGenerator = SyntaxGenerator.GetGenerator(document);
 
-        SyntaxNode fixedMethodDeclarationNode = syntaxGenerator.MethodDeclaration(
-            methodSymbol.Name,
-            GetParameters(syntaxGenerator, isParameterLess, wellKnownTypeProvider),
-            typeParameters: null,
-            GetReturnType(syntaxGenerator, methodSymbol, wellKnownTypeProvider),
-            Accessibility.Public,
-            GetModifiers(methodSymbol, shouldBeStatic),
-            GetStatements(node, syntaxGenerator));
+        // Start with the original method declaration to preserve trivia and body
+        MethodDeclarationSyntax fixedMethodDeclaration = methodDeclaration;
 
-        // Copy the attributes from the old method to the new method.
-        fixedMethodDeclarationNode = syntaxGenerator.AddAttributes(fixedMethodDeclarationNode, syntaxGenerator.GetAttributes(node));
+        // Update parameters
+        ParameterListSyntax newParameterList = GetParameterList(isParameterLess, wellKnownTypeProvider);
+        fixedMethodDeclaration = fixedMethodDeclaration.WithParameterList(newParameterList);
 
-        return document.WithSyntaxRoot(root.ReplaceNode(node, fixedMethodDeclarationNode)).Project.Solution;
-    }
+        // Update return type
+        TypeSyntax? newReturnType = GetReturnType(methodSymbol, wellKnownTypeProvider);
+        if (newReturnType is not null)
+        {
+            fixedMethodDeclaration = fixedMethodDeclaration.WithReturnType(newReturnType);
+        }
+
+        // Update modifiers (accessibility and static)
+        SyntaxTokenList newModifiers = GetModifiers(methodDeclaration, methodSymbol, shouldBeStatic);
+        fixedMethodDeclaration = fixedMethodDeclaration.WithModifiers(newModifiers);
+
+        // Remove type parameters if any
+        if (fixedMethodDeclaration.TypeParameterList is not null)
+        {
+            fixedMethodDeclaration = fixedMethodDeclaration.WithTypeParameterList(null);
+        }
+
+        // Remove return and yield return statements from body if needed
+        if (fixedMethodDeclaration.Body is not null)
+        {
+            SyntaxList<StatementSyntax> statements = fixedMethodDeclaration.Body.Statements;
+            IEnumerable<StatementSyntax> filteredStatements = statements
+                .Where(x => !x.IsKind(SyntaxKind.ReturnStatement) && !x.IsKind(SyntaxKind.YieldReturnStatement));
+
+            if (statements.Count != filteredStatements.Count())
+            {
+                fixedMethodDeclaration = fixedMethodDeclaration.WithBody(
+                    fixedMethodDeclaration.Body.WithStatements(SyntaxFactory.List(filteredStatements)));
+            }
+        }
 
-    private static IEnumerable<SyntaxNode> GetStatements(SyntaxNode node, SyntaxGenerator syntaxGenerator)
-        => syntaxGenerator.GetStatements(node)
-            .Where(x => !x.IsKind(SyntaxKind.ReturnStatement) && !x.IsKind(SyntaxKind.YieldReturnStatement));
+        return document.WithSyntaxRoot(root.ReplaceNode(node, fixedMethodDeclaration)).Project.Solution;
+    }
 
-    private static DeclarationModifiers GetModifiers(IMethodSymbol methodSymbol, bool shouldBeStatic)
+    private static SyntaxTokenList GetModifiers(MethodDeclarationSyntax methodDeclaration, IMethodSymbol methodSymbol, bool shouldBeStatic)
     {
-        DeclarationModifiers newModifiers = methodSymbol.IsAsync
-            ? DeclarationModifiers.Async
-            : DeclarationModifiers.None;
+        // Start with existing modifiers
+        SyntaxTokenList modifiers = methodDeclaration.Modifiers;
+
+        // Remove all accessibility modifiers
+        modifiers = SyntaxFactory.TokenList(modifiers.Where(m =>
+            !m.IsKind(SyntaxKind.PublicKeyword) &&
+            !m.IsKind(SyntaxKind.PrivateKeyword) &&
+            !m.IsKind(SyntaxKind.ProtectedKeyword) &&
+            !m.IsKind(SyntaxKind.InternalKeyword)));
+
+        // Remove abstract modifier if present
+        modifiers = SyntaxFactory.TokenList(modifiers.Where(m => !m.IsKind(SyntaxKind.AbstractKeyword)));
+
+        // Handle static modifier
+        bool hasStaticModifier = modifiers.Any(m => m.IsKind(SyntaxKind.StaticKeyword));
+        if (shouldBeStatic && !hasStaticModifier)
+        {
+            modifiers = modifiers.Add(SyntaxFactory.Token(SyntaxKind.StaticKeyword));
+        }
+        else if (!shouldBeStatic && hasStaticModifier)
+        {
+            modifiers = SyntaxFactory.TokenList(modifiers.Where(m => !m.IsKind(SyntaxKind.StaticKeyword)));
+        }
 
-        return newModifiers.WithIsStatic(shouldBeStatic);
+        // Add public at the beginning
+        modifiers = SyntaxFactory.TokenList(SyntaxFactory.Token(SyntaxKind.PublicKeyword)).AddRange(modifiers);
+
+        return modifiers;
     }
 
-    private static SyntaxNode? GetReturnType(SyntaxGenerator syntaxGenerator, IMethodSymbol methodSymbol, WellKnownTypeProvider wellKnownTypeProvider)
+    private static TypeSyntax? GetReturnType(IMethodSymbol methodSymbol, WellKnownTypeProvider wellKnownTypeProvider)
     {
         if (SymbolEqualityComparer.Default.Equals(methodSymbol.ReturnType.OriginalDefinition, wellKnownTypeProvider.GetOrCreateTypeByMetadataName(WellKnownTypeNames.SystemThreadingTasksValueTask1)))
         {
-            return syntaxGenerator.IdentifierName("ValueTask");
+            return SyntaxFactory.IdentifierName("ValueTask");
         }
 
         if (methodSymbol.IsAsync
             || SymbolEqualityComparer.Default.Equals(methodSymbol.ReturnType.OriginalDefinition, wellKnownTypeProvider.GetOrCreateTypeByMetadataName(WellKnownTypeNames.SystemThreadingTasksTask1)))
         {
-            return syntaxGenerator.IdentifierName("Task");
+            return SyntaxFactory.IdentifierName("Task");
         }
 
-        // For all other cases return void.
-        return VoidReturnTypeNode;
+        // For void, return a predefined void type
+        return SyntaxFactory.PredefinedType(SyntaxFactory.Token(SyntaxKind.VoidKeyword));
     }
 
-    private static IEnumerable<SyntaxNode> GetParameters(SyntaxGenerator syntaxGenerator, bool isParameterLess,
-        WellKnownTypeProvider wellKnownTypeProvider)
+    private static ParameterListSyntax GetParameterList(bool isParameterLess, WellKnownTypeProvider wellKnownTypeProvider)
     {
         if (isParameterLess
             || !wellKnownTypeProvider.TryGetOrCreateTypeByMetadataName(
                 WellKnownTypeNames.MicrosoftVisualStudioTestToolsUnitTestingTestContext,
                 out INamedTypeSymbol? testContextTypeSymbol))
         {
-            return [];
+            return SyntaxFactory.ParameterList();
         }
 
-        SyntaxNode testContextType = syntaxGenerator.TypeExpression(testContextTypeSymbol);
-        SyntaxNode testContextParameter = syntaxGenerator.ParameterDeclaration("testContext", testContextType);
-        return [testContextParameter];
+        TypeSyntax testContextType = SyntaxFactory.ParseTypeName(testContextTypeSymbol.ToDisplayString());
+        ParameterSyntax testContextParameter = SyntaxFactory.Parameter(SyntaxFactory.Identifier("testContext"))
+            .WithType(testContextType);
+
+        return SyntaxFactory.ParameterList(SyntaxFactory.SingletonSeparatedList(testContextParameter));
     }
 }

test/UnitTests/MSTest.Analyzers.UnitTests/ClassInitializeShouldBeValidAnalyzerTests.cs

@@ -711,4 +711,51 @@ public static void ClassInitialize(TestContext testContext)
 
         await VerifyCS.VerifyAnalyzerAsync(code);
     }
+
+    [TestMethod]
+    public async Task WhenClassInitializeHasComments_CommentsArePreserved()
+    {
+        string code = """
+            using Microsoft.VisualStudio.TestTools.UnitTesting;
+
+            [TestClass]
+            public class MyTestClass
+            {
+                [ClassInitialize]
+                internal static void {|#0:ClassInitialize|}(TestContext testContext)
+                {
+                    InitializeClass();
+
+                    // Class initialization comments;
+                    // Setup code here
+                }
+
+                private static void InitializeClass() { }
+            }
+            """;
+
+        string fixedCode = """
+            using Microsoft.VisualStudio.TestTools.UnitTesting;
+
+            [TestClass]
+            public class MyTestClass
+            {
+                [ClassInitialize]
+                public static void ClassInitialize(TestContext testContext)
+                {
+                    InitializeClass();
+
+                    // Class initialization comments;
+                    // Setup code here
+                }
+
+                private static void InitializeClass() { }
+            }
+            """;
+
+        await VerifyCS.VerifyCodeFixAsync(
+            code,
+            VerifyCS.Diagnostic().WithLocation(0).WithArguments("ClassInitialize"),
+            fixedCode);
+    }
 }

test/UnitTests/MSTest.Analyzers.UnitTests/TestCleanupShouldBeValidAnalyzerTests.cs

@@ -552,4 +552,51 @@ public void TestCleanup()
 
         await VerifyCS.VerifyAnalyzerAsync(code);
     }
+
+    [TestMethod]
+    public async Task WhenTestCleanupHasComments_CommentsArePreserved()
+    {
+        string code = """
+            using Microsoft.VisualStudio.TestTools.UnitTesting;
+
+            [TestClass]
+            public class MyTestClass
+            {
+                [TestCleanup]
+                internal void {|#0:TestCleanup|}()
+                {
+                    CleanupCode();
+
+                    // Cleanup comments;
+                    // More comments
+                }
+
+                private void CleanupCode() { }
+            }
+            """;
+
+        string fixedCode = """
+            using Microsoft.VisualStudio.TestTools.UnitTesting;
+
+            [TestClass]
+            public class MyTestClass
+            {
+                [TestCleanup]
+                public void TestCleanup()
+                {
+                    CleanupCode();
+
+                    // Cleanup comments;
+                    // More comments
+                }
+
+                private void CleanupCode() { }
+            }
+            """;
+
+        await VerifyCS.VerifyCodeFixAsync(
+            code,
+            VerifyCS.Diagnostic().WithLocation(0).WithArguments("TestCleanup"),
+            fixedCode);
+    }
 }

test/UnitTests/MSTest.Analyzers.UnitTests/TestInitializeShouldBeValidAnalyzerTests.cs

@@ -478,6 +478,51 @@ await VerifyCS.VerifyCodeFixAsync(
             fixedCode);
     }
 
+    [TestMethod]
+    public async Task WhenTestInitializeHasComments_CommentsArePreserved()
+    {
+        string code = """
+            using Microsoft.VisualStudio.TestTools.UnitTesting;
+
+            [TestClass]
+            public class MyTestClass
+            {
+                [TestInitialize]
+                public void {|#0:TestSetup|}(TestContext tc)
+                {
+                    SomeCode();
+
+                    // Some comments;
+                }
+
+                private void SomeCode() { }
+            }
+            """;
+
+        string fixedCode = """
+            using Microsoft.VisualStudio.TestTools.UnitTesting;
+
+            [TestClass]
+            public class MyTestClass
+            {
+                [TestInitialize]
+                public void TestSetup()
+                {
+                    SomeCode();
+
+                    // Some comments;
+                }
+
+                private void SomeCode() { }
+            }
+            """;
+
+        await VerifyCS.VerifyCodeFixAsync(
+            code,
+            VerifyCS.Diagnostic().WithLocation(0).WithArguments("TestSetup"),
+            fixedCode);
+    }
+
     [TestMethod]
     public async Task WhenTestInitializeIsNotOnClass_Diagnostic()
     {