Fix Is.EqualTo for collections (#16)

meziantou · Dec 4, 2022 · af216ac · af216ac
1 parent 6d28640
commit af216ac
Show file tree

Hide file tree

Showing 2 changed files with 104 additions and 2 deletions.
diff --git a/Meziantou.FluentAssertionsAnalyzers.Tests/NUnitToFluentAssertionsAnalyzerUnitTests.cs b/Meziantou.FluentAssertionsAnalyzers.Tests/NUnitToFluentAssertionsAnalyzerUnitTests.cs
@@ -52,6 +52,72 @@ public void MyTest()
 """);
     }
 
+    [Fact]
+    public Task EnumerableTest()
+    {
+        return Assert(
+            $$"""
+using System.Collections;
+using NUnit.Framework;
+
+class CustomEnumerable : IEnumerable
+{
+    private readonly IEnumerable m_Collection;
+
+    public CustomEnumerable(IEnumerable collection)
+    {
+        m_Collection = collection;
+    }
+    public IEnumerator GetEnumerator()
+    {
+        foreach (var item in m_Collection)
+        {
+            yield return item;
+        }
+    }
+}
+
+class Test
+{
+    public void MyTest()
+    {
+        var collection = new CustomEnumerable(new int[0]);
+        [||]Assert.That(collection, Is.EqualTo(new int[0]));
+    }
+}
+""",
+            $$"""
+using System.Collections;
+using NUnit.Framework;
+
+class CustomEnumerable : IEnumerable
+{
+    private readonly IEnumerable m_Collection;
+
+    public CustomEnumerable(IEnumerable collection)
+    {
+        m_Collection = collection;
+    }
+    public IEnumerator GetEnumerator()
+    {
+        foreach (var item in m_Collection)
+        {
+            yield return item;
+        }
+    }
+}
+
+class Test
+{
+    public void MyTest()
+    {
+        var collection = new CustomEnumerable(new int[0]);
+        Assert.That(collection, Is.EqualTo(new int[0]));
+    }
+}
+""");
+    }
+
     [Fact]
     public Task Assert_Dynamic()
     {
@@ -614,6 +680,9 @@ public void MyTest()
 
     [InlineData(@"Assert.That("""", Is.EqualTo(""expected""))", @""""".Should().Be(""expected"")")]
     [InlineData(@"Assert.That("""", Is.Not.EqualTo(""expected""))", @""""".Should().NotBe(""expected"")")]
+    [InlineData(@"Assert.That(collection, Is.EqualTo(new int[0]))", @"collection.Should().Equal(new int[0])")]
+    [InlineData(@"Assert.That(collection, Is.Not.EqualTo(new int[0]))", @"collection.Should().NotEqual(new int[0])")]
+    [InlineData(@"Assert.That((IEnumerable<int>)collection, Is.EqualTo(new int[0]))", @"((IEnumerable<int>)collection).Should().Equal(new int[0])")]
 
     [InlineData(@"Assert.That(collection, Is.EquivalentTo(new int[0]))", @"collection.Should().BeEquivalentTo(new int[0])")]
     [InlineData(@"Assert.That(collection, Is.Not.EquivalentTo(new int[0]))", @"collection.Should().NotBeEquivalentTo(new int[0])")]
@@ -649,6 +718,7 @@ public Task Assert_Tests(string code, string fix)
     {
         return Assert(
 $$"""
+using System.Collections.Generic;
 using NUnit.Framework;
 
 class Test
@@ -661,6 +731,7 @@ public void MyTest()
 }
 """,
 $$"""
+using System.Collections.Generic;
 using FluentAssertions;
 using NUnit.Framework;
 

diff --git a/Meziantou.FluentAssertionsAnalyzers/NunitAssertAnalyzerCodeFixProvider.cs b/Meziantou.FluentAssertionsAnalyzers/NunitAssertAnalyzerCodeFixProvider.cs
@@ -478,11 +478,21 @@ private static async Task<Document> Rewrite(Document document, SyntaxNode nodeTo
                         }
                         else if (IsMethod(out var expected, isSymbol, "EqualTo"))
                         {
-                            result = rewrite.UsingShould(arguments[0], "Be", ArgumentList(expected, arguments.Skip(2)));
+                            var (isCollection, isMigrationSupported) = IsCollection(arguments[0]);
+                            if (isMigrationSupported)
+                            {
+                                var replacementMethodName = isCollection ? "Equal": "Be";
+                                result = rewrite.UsingShould(arguments[0], replacementMethodName, ArgumentList(expected, arguments.Skip(2)));
+                            }
                         }
                         else if (IsMethod(out expected, isSymbol, "Not", "EqualTo"))
                         {
-                            result = rewrite.UsingShould(arguments[0], "NotBe", ArgumentList(expected, arguments.Skip(2)));
+                            var (isCollection, isMigrationSupported) = IsCollection(arguments[0]);
+                            if (isMigrationSupported)
+                            {
+                                var replacementMethodName = isCollection ? "NotEqual" : "NotBe";
+                                result = rewrite.UsingShould(arguments[0], replacementMethodName, ArgumentList(expected, arguments.Skip(2)));
+                            }
                         }
                         else if (IsMethod(out expected, isSymbol, "SameAs"))
                         {
@@ -582,6 +592,27 @@ private static async Task<Document> Rewrite(Document document, SyntaxNode nodeTo
                         }
                     }
 
+                    (bool isCollection, bool isMigrationSupported) IsCollection(ArgumentSyntax argumentSyntax)
+                    {
+                        var argumentTypeSymbol = semanticModel.GetTypeInfo(argumentSyntax.Expression, cancellationToken).Type;
+                        if (argumentTypeSymbol == null || argumentTypeSymbol.SpecialType == SpecialType.System_String)
+                            return (false, true);
+
+                        var isCollection = 
+                               argumentTypeSymbol.OriginalDefinition.SpecialType == SpecialType.System_Collections_IEnumerable
+                            || argumentTypeSymbol.OriginalDefinition.AllInterfaces.Any(i => 
+                                i.SpecialType == SpecialType.System_Collections_IEnumerable);
+
+                        var isSupportedCollection = argumentTypeSymbol.OriginalDefinition.TypeKind == TypeKind.Array
+                            || argumentTypeSymbol.OriginalDefinition.SpecialType == SpecialType.System_Collections_Generic_IEnumerable_T
+                            || argumentTypeSymbol.OriginalDefinition.AllInterfaces.Any(i => 
+                                i.SpecialType == SpecialType.System_Collections_Generic_IEnumerable_T);
+
+                        var isMigrationSupported = !isCollection || isSupportedCollection;
+
+                        return (isCollection, isMigrationSupported);
+                    }
+
                     bool Is(ITypeSymbol root, params string[] memberNames)
                     {
                         var currentOp = op;