Keep leading and trailing trivia when replacing statements with Assert.Multiple

src/nunit.analyzers.tests/UseAssertMultiple/UseAssertMultipleCodeFixTests.cs

@@ -134,5 +134,63 @@ private sealed class Configuration
         }");
             RoslynAssert.FixAll(analyzer, fix, expectedDiagnostic, code, fixedCode);
         }
+
+        [TestCase("            // ", "")]
+        [TestCase("\r\n            // ", "")]
+        [TestCase("\r\n            ", "")]
+        [TestCase("\r\n            ", " // Same line Comment")]
+        [TestCase("\r\n            ", "\r\n            // Final Comment on next line")]
+        public void VerifyKeepsTrivia(string separation, string comment)
+        {
+            var code = TestUtility.WrapMethodInClassNamespaceAndAddUsings(@$"
+        public void TestMethod()
+        {{
+            const bool True = true;
+            const bool False = false;
+
+            // Verify that our bool constants are correct
+            ↓Assert.That(True, Is.True);
+            Assert.That(False, Is.False);
+{separation}Console.WriteLine(""Next Statement"");{comment}
+        }}");
+            var fixedCode = TestUtility.WrapMethodInClassNamespaceAndAddUsings(@$"
+        public void TestMethod()
+        {{
+            const bool True = true;
+            const bool False = false;
+
+            Assert.Multiple(() =>
+            {{
+                // Verify that our bool constants are correct
+                Assert.That(True, Is.True);
+                Assert.That(False, Is.False);
+            }});
+{separation}Console.WriteLine(""Next Statement"");{comment}
+        }}");
+            RoslynAssert.CodeFix(analyzer, fix, expectedDiagnostic, code, fixedCode);
+        }
+
+        [Test]
+        public void VerifyKeepsTrivia()
+        {
+            var code = TestUtility.WrapMethodInClassNamespaceAndAddUsings(@$"
+        public void TestMethod()
+        {{
+                // Verify that boolean work as expected
+            ↓Assert.That(true, Is.True);
+            Assert.That(false, Is.False);
+        }}");
+            var fixedCode = TestUtility.WrapMethodInClassNamespaceAndAddUsings(@$"
+        public void TestMethod()
+        {{
+            Assert.Multiple(() =>
+            {{
+                // Verify that boolean work as expected
+                Assert.That(true, Is.True);
+                Assert.That(false, Is.False);
+            }});
+        }}");
+            RoslynAssert.CodeFix(analyzer, fix, expectedDiagnostic, code, fixedCode);
+        }
     }
 }

src/nunit.analyzers/UseAssertMultiple/UseAssertMultipleCodeFix.cs

@@ -91,6 +91,22 @@ public override async Task RegisterCodeFixesAsync(CodeFixContext context)
                 statements[when].Add(statement);
             }
 
+            // If there was an empty line between the code above the Assert, then keep it.
+            SyntaxTrivia? endOfLineTrivia = default;
+
+            var firstStatement = statementsInsideAssertMultiple[0];
+            if (firstStatement.HasLeadingTrivia)
+            {
+                SyntaxTriviaList trivia = firstStatement.GetLeadingTrivia();
+                SyntaxTrivia firstTrivia = trivia.First();
+                if (firstTrivia.IsKind(SyntaxKind.EndOfLineTrivia))
+                {
+                    // Remember the trivia and delete it from the first statement inside the Assert.Multiple
+                    endOfLineTrivia = firstTrivia;
+                    statementsInsideAssertMultiple[0] = firstStatement.ReplaceTrivia(firstTrivia, Enumerable.Empty<SyntaxTrivia>());
+                }
+            }
+
             ParenthesizedLambdaExpressionSyntax parenthesizedLambdaExpression =
                 SyntaxFactory.ParenthesizedLambdaExpression(
                     SyntaxFactory.Block(statementsInsideAssertMultiple));
@@ -113,8 +129,21 @@ public override async Task RegisterCodeFixesAsync(CodeFixContext context)
                             SyntaxFactory.Argument(parenthesizedLambdaExpression)
                         })))).WithAdditionalAnnotations(Formatter.Annotation);
 
+            if (endOfLineTrivia is not null)
+            {
+                // Add the remembered blank line.
+                assertMultiple = assertMultiple.WithLeadingTrivia(endOfLineTrivia.Value);
+            }
+
+            // Add new line after the Assert.Multiple statement.
+            assertMultiple = assertMultiple.WithTrailingTrivia(SyntaxFactory.CarriageReturnLineFeed);
+
+            // Comments at the end of a block are not associated with the last statement but with the closing brace
+            // Keep the exising block's open and close braces with associated trivia in our updated block.
             var updatedBlock = SyntaxFactory.Block(
-                statementsBeforeAssertMultiple.Append(assertMultiple).Concat(statementsAfterAssertMultiple));
+                block.OpenBraceToken,
+                SyntaxFactory.List(statementsBeforeAssertMultiple.Append(assertMultiple).Concat(statementsAfterAssertMultiple)),
+                block.CloseBraceToken);
 
             SyntaxNode newRoot = root.ReplaceNode(block, updatedBlock);