Fix for already async method

thomhurst · Feb 1, 2025 · 215d845 · 215d845
1 parent 713b3d2
commit 215d845
Show file tree

Hide file tree

Showing 2 changed files with 75 additions and 31 deletions.
diff --git a/TUnit.Assertions.Analyzers.CodeFixers.Tests/AwaitAssertionCodeFixProviderTests.cs b/TUnit.Assertions.Analyzers.CodeFixers.Tests/AwaitAssertionCodeFixProviderTests.cs
@@ -43,4 +43,44 @@ public async Task MyTest()
                 """
             );
     }
+
+    [Test]
+    public async Task Async_Task_Just_Awaits()
+    {
+        await Verifier
+            .VerifyCodeFixAsync(
+                """
+                using System.Threading.Tasks;
+                using TUnit.Assertions;
+                using TUnit.Assertions.Extensions;
+                using TUnit.Core;
+
+                public class MyClass
+                {
+                    public async Task MyTest()
+                    {
+                        await Task.Delay(1);
+                        {|#0:Assert.That(1)|}.IsEqualTo(1);
+                    }
+                }
+                """,
+                Verifier.Diagnostic(Rules.AwaitAssertion)
+                    .WithLocation(0),
+                """
+                using System.Threading.Tasks;
+                using TUnit.Assertions;
+                using TUnit.Assertions.Extensions;
+                using TUnit.Core;
+
+                public class MyClass
+                {
+                    public async Task MyTest()
+                    {
+                        await Task.Delay(1);
+                        await Assert.That(1).IsEqualTo(1);
+                    }
+                }
+                """
+            );
+    }
 }
diff --git a/TUnit.Assertions.Analyzers.CodeFixers/AwaitAssertionCodeFixProvider.cs b/TUnit.Assertions.Analyzers.CodeFixers/AwaitAssertionCodeFixProvider.cs
@@ -68,42 +68,46 @@ private static async Task<Document> AwaitAssertionAsync(Document document, Expre
             .OfType<MethodDeclarationSyntax>()
             .FirstOrDefault();
 
-        if (methodDeclaration != null)
+        if (methodDeclaration == null)
         {
-            // Check if the method is already async
-            if (!methodDeclaration.Modifiers.Any(SyntaxKind.AsyncKeyword))
+            return editor.GetChangedDocument();
+        }
+
+        var modifiers = methodDeclaration.Modifiers;
+
+        var returnType = methodDeclaration.ReturnType;
+        var newReturnType = returnType;
+
+        // Check if the method is already async
+        if (!methodDeclaration.Modifiers.Any(SyntaxKind.AsyncKeyword))
+        {
+            // Add async modifier
+            var asyncModifier = SyntaxFactory.Token(SyntaxKind.AsyncKeyword);
+            modifiers = methodDeclaration.Modifiers.Add(asyncModifier
+                .WithTrailingTrivia(SyntaxFactory.Space));
+
+            // Update the return type to Task or Task<T>
+            if (returnType is PredefinedTypeSyntax predefinedType &&
+                predefinedType.Keyword.IsKind(SyntaxKind.VoidKeyword))
             {
-                // Add async modifier
-                var asyncModifier = SyntaxFactory.Token(SyntaxKind.AsyncKeyword);
-                var newModifiers = methodDeclaration.Modifiers.Add(asyncModifier
-                    .WithTrailingTrivia(SyntaxFactory.Space));
-
-                // Update the return type to Task or Task<T>
-                var returnType = methodDeclaration.ReturnType;
-                var newReturnType = returnType;
-
-                if (returnType is PredefinedTypeSyntax predefinedType &&
-                    predefinedType.Keyword.IsKind(SyntaxKind.VoidKeyword))
-                {
-                    newReturnType = SyntaxFactory.IdentifierName("Task")
-                        .WithTrailingTrivia(SyntaxFactory.Space);
-                }
-                else if (returnType is not GenericNameSyntax genericName || genericName.Identifier.Text != "Task")
-                {
-                    newReturnType = SyntaxFactory.ParseTypeName($"Task<{returnType}>")
-                        .WithTrailingTrivia(SyntaxFactory.Space);
-                }
-
-                var newMethodDeclaration = methodDeclaration
-                    .ReplaceNode(expressionSyntax, awaitExpression)
-                    .WithModifiers(newModifiers)
-                    .WithReturnType(newReturnType)
-                    .WithAdditionalAnnotations(Formatter.Annotation);
-
-                editor.ReplaceNode(methodDeclaration, newMethodDeclaration);
+                newReturnType = SyntaxFactory.IdentifierName("Task")
+                    .WithTrailingTrivia(SyntaxFactory.Space);
+            }
+            else if (returnType is not GenericNameSyntax genericName || genericName.Identifier.Text != "Task")
+            {
+                newReturnType = SyntaxFactory.ParseTypeName($"Task<{returnType}>")
+                    .WithTrailingTrivia(SyntaxFactory.Space);
             }
         }
 
+        var newMethodDeclaration = methodDeclaration
+            .ReplaceNode(expressionSyntax, awaitExpression)
+            .WithModifiers(modifiers)
+            .WithReturnType(newReturnType)
+            .WithAdditionalAnnotations(Formatter.Annotation);
+
+        editor.ReplaceNode(methodDeclaration, newMethodDeclaration);
+
         return editor.GetChangedDocument();
     }
 }