Improve support for XCTAssertEqual and XCTAssertNil in noForceUnwrapInTests (#2208)

Rules.md

@@ -1684,12 +1684,12 @@ Use XCTUnwrap or #require in test cases, rather than force unwrapping.
 -           let myValue = foo.bar!.value as! Value
 -           let otherValue = (foo! as! Other).bar
 -           otherValue.manager!.prepare()
--           #expect(myValue.property! == other)
+-           #expect(myValue!.property! == other)
 +       @Test func myFeature() throws {
 +           let myValue = try #require(foo.bar?.value as? Value)
 +           let otherValue = try #require((foo as? Other)?.bar)
 +           otherValue.manager?.prepare()
-+           #expect(try #require(myValue.property) == other)
++           #expect(myValue?.property == other)
       }
     }
 
@@ -1699,11 +1699,11 @@ Use XCTUnwrap or #require in test cases, rather than force unwrapping.
 -       func testMyFeature() {
 -           let myValue = foo.bar!.value as! Value
 -           let otherValue = (foo! as! Other).bar
--           XCTAssertEqual(myValue.property, "foo")
+-           XCTAssertEqual(myValue!.property!, "foo")
 +       func testMyFeature() throws {
 +           let myValue = try XCTUnwrap(foo.bar?.value as? Value)
 +           let otherValue = try XCTUnwrap((foo as? Other)?.bar)
-+           XCTAssertEqual(try XCTUnwrap(myValue.property), otherValue)
++           XCTAssertEqual(myValue?.property, otherValue)
       }
     }
 ```

Sources/Rules/NoForceUnwrapInTests.swift

@@ -154,6 +154,34 @@ public extension FormatRule {
                     needsUnwrapMethod = false
                 }
 
+                // If this expression is followed by ==, changing `foo!.bar == bar` to `foo?.bar == bar` is a safe change as-is
+                if let tokenAfterExpression = formatter.index(of: .nonSpaceOrCommentOrLinebreak, after: expressionRange.upperBound),
+                   formatter.tokens[tokenAfterExpression] == .operator("==", .infix)
+                {
+                    needsUnwrapMethod = false
+                }
+
+                // If this expression is within XCTAssertEqual or XCTAssertNil, changing `foo!.bar` to `foo?.bar` is a safe change as-is,
+                // as long as this isn't a subexpression within a parent operator expression.
+                if let containingParenScope = formatter.startOfScope(at: expressionRange.lowerBound),
+                   formatter.tokens[containingParenScope] == .startOfScope("("),
+                   let functionNameIndex = formatter.index(of: .nonSpaceOrCommentOrLinebreak, before: containingParenScope),
+                   formatter.tokens[functionNameIndex].isIdentifier,
+                   let tokenAfterExpression = formatter.index(of: .nonSpaceOrCommentOrLinebreak, after: expressionRange.upperBound),
+                   !formatter.tokens[tokenAfterExpression].isOperator
+                {
+                    let functionName = formatter.tokens[functionNameIndex].string
+                    if functionName == "XCTAssertNil" {
+                        needsUnwrapMethod = false
+                    } else if functionName == "XCTAssertEqual" {
+                        // Ensure this is `XCTAssertEqual(_:_:)`, not `XCTAssertEqual(_:_:accuracy:)` (which doesn't support optionals)
+                        let arguments = formatter.parseFunctionCallArguments(startOfScope: containingParenScope)
+                        if arguments.count == 2 {
+                            needsUnwrapMethod = false
+                        }
+                    }
+                }
+
                 // If this expression is a standalone method call like `foo!.bar()`, then `foo?.bar()` works perfectly well.
                 // Heuristic: If the scope containing this code is a code block, and the previous token is part of a completely
                 // separate expression (or, the start of the function body), then this is a standalone expression.
@@ -218,12 +246,12 @@ public extension FormatRule {
         -           let myValue = foo.bar!.value as! Value
         -           let otherValue = (foo! as! Other).bar
         -           otherValue.manager!.prepare()
-        -           #expect(myValue.property! == other)
+        -           #expect(myValue!.property! == other)
         +       @Test func myFeature() throws {
         +           let myValue = try #require(foo.bar?.value as? Value)
         +           let otherValue = try #require((foo as? Other)?.bar)
         +           otherValue.manager?.prepare()
-        +           #expect(try #require(myValue.property) == other)
+        +           #expect(myValue?.property == other)
               }
             }
 
@@ -233,11 +261,11 @@ public extension FormatRule {
         -       func testMyFeature() {
         -           let myValue = foo.bar!.value as! Value
         -           let otherValue = (foo! as! Other).bar
-        -           XCTAssertEqual(myValue.property, "foo")
+        -           XCTAssertEqual(myValue!.property!, "foo")
         +       func testMyFeature() throws {
         +           let myValue = try XCTUnwrap(foo.bar?.value as? Value)
         +           let otherValue = try XCTUnwrap((foo as? Other)?.bar)
-        +           XCTAssertEqual(try XCTUnwrap(myValue.property), otherValue)
+        +           XCTAssertEqual(myValue?.property, otherValue)
               }
             }
         ```

Tests/CodeOrganizationTests.swift

@@ -106,7 +106,9 @@ class CodeOrganizationTests: XCTestCase {
                 // If this is a function call, parse the labels to disambiguate
                 // between methods with the same base name
                 var functionCallArguments: [String?]?
-                if let functionCallStartOfScope = formatter.index(of: .startOfScope("("), after: index) {
+                if let functionCallStartOfScope = formatter.index(of: .nonSpaceOrCommentOrLinebreak, after: index),
+                   formatter.tokens[functionCallStartOfScope] == .startOfScope("(")
+                {
                     functionCallArguments = formatter.parseFunctionCallArguments(startOfScope: functionCallStartOfScope).map(\.label)
                 }
 

Tests/Rules/NoForceUnwrapInTestsTests.swift

@@ -111,7 +111,7 @@ final class NoForceUnwrapInTestsTests: XCTestCase {
         class TestCase: XCTestCase {
             func test_functionCall() throws {
                 someFunction(try XCTUnwrap(myOptional), try XCTUnwrap(anotherOptional))
-                XCTAssertEqual(try XCTUnwrap(result?.property), "expected")
+                XCTAssertEqual(result?.property, "expected")
             }
         }
         """
@@ -140,7 +140,7 @@ final class NoForceUnwrapInTestsTests: XCTestCase {
             func test_ifStatement() throws {
                 if
                     try XCTUnwrap(foo?.bar()),
-                    try XCTUnwrap(myOptional?.value) == someValue
+                    myOptional?.value == someValue
                 {
                     // do something
                 }
@@ -198,7 +198,7 @@ final class NoForceUnwrapInTestsTests: XCTestCase {
             func test_guardStatement() throws {
                 guard
                     try XCTUnwrap(foo?.bar()),
-                    try XCTUnwrap(myOptional?.value) == someValue
+                    myOptional?.value == someValue
                 else {
                     return
                 }
@@ -420,7 +420,7 @@ final class NoForceUnwrapInTestsTests: XCTestCase {
         class TestCase: XCTestCase {
             func test_complexExpression() throws {
                 XCTAssertEqual(
-                    try XCTUnwrap(myDictionary["key"]?.processedValue(with: try XCTUnwrap(parameter))),
+                    myDictionary["key"]?.processedValue(with: try XCTUnwrap(parameter)),
                     expectedResult
                 )
             }
@@ -506,7 +506,7 @@ final class NoForceUnwrapInTestsTests: XCTestCase {
 
         class TestCase: XCTestCase {
             func test_forceCasts() throws {
-                XCTAssertEqual(try XCTUnwrap(route.query as? [String: String]), ["a": "b"])
+                XCTAssertEqual(route.query as? [String: String], ["a": "b"])
                 XCTAssert(try XCTUnwrap((foo as? Bar)?.baaz))
                 XCTAssert(try XCTUnwrap((foo as? Bar)?.baaz))
             }
@@ -659,4 +659,124 @@ final class NoForceUnwrapInTestsTests: XCTestCase {
         """
         testFormatting(for: input, output, rule: .noForceUnwrapInTests)
     }
+
+    func testXCTAssertEqual_KeepsForceUnwrapsAsOptionalChaining() throws {
+        let input = """
+        import XCTest
+
+        class TestCase: XCTestCase {
+            func test_something() {
+                XCTAssertEqual(foo!.bar, baaz!.quux)
+            }
+        }
+        """
+        let output = """
+        import XCTest
+
+        class TestCase: XCTestCase {
+            func test_something() {
+                XCTAssertEqual(foo?.bar, baaz?.quux)
+            }
+        }
+        """
+        testFormatting(for: input, output, rule: .noForceUnwrapInTests)
+    }
+
+    func testXCTAssertNil_KeepsForceUnwrapsAsOptionalChaining() throws {
+        let input = """
+        import XCTest
+
+        class TestCase: XCTestCase {
+            func test_something() {
+                XCTAssertNil(foo!.bar)
+            }
+        }
+        """
+        let output = """
+        import XCTest
+
+        class TestCase: XCTestCase {
+            func test_something() {
+                XCTAssertNil(foo?.bar)
+            }
+        }
+        """
+        testFormatting(for: input, output, rule: .noForceUnwrapInTests)
+    }
+
+    func testEqualityComparison_KeepsForceUnwrapsAsOptionalChaining() throws {
+        let input = """
+        import Testing
+
+        @Test func something() {
+            #expect(foo!.bar == baaz)
+        }
+        """
+        let output = """
+        import Testing
+
+        @Test func something() {
+            #expect(foo?.bar == baaz)
+        }
+        """
+        testFormatting(for: input, output, rule: .noForceUnwrapInTests)
+    }
+
+    func testEqualityComparisonWithNil_KeepsForceUnwrapsAsOptionalChaining() throws {
+        let input = """
+        import Testing
+
+        @Test func something() {
+            #expect(foo!.bar == nil)
+        }
+        """
+        let output = """
+        import Testing
+
+        @Test func something() {
+            #expect(foo?.bar == nil)
+        }
+        """
+        testFormatting(for: input, output, rule: .noForceUnwrapInTests)
+    }
+
+    func testXCTAssertEqualWithAccuracy_RequiresXCTUnwrap() throws {
+        let input = """
+        import XCTest
+
+        class TestCase: XCTestCase {
+            func test_something() {
+                XCTAssertEqual(foo!.value, 3.14, accuracy: 0.01)
+            }
+        }
+        """
+        let output = """
+        import XCTest
+
+        class TestCase: XCTestCase {
+            func test_something() throws {
+                XCTAssertEqual(try XCTUnwrap(foo?.value), 3.14, accuracy: 0.01)
+            }
+        }
+        """
+        testFormatting(for: input, output, rule: .noForceUnwrapInTests)
+    }
+
+    func testForceUnwrapWithOperatorFollowing_RequiresXCTUnwrap() throws {
+        let input = """
+        import Testing
+
+        @Test func something() {
+            #expect(foo!.bar + 2 == 3)
+        }
+        """
+        let output = """
+        import Testing
+
+        @Test func something() throws {
+            #expect(try #require(foo?.bar) + 2 == 3)
+        }
+        """
+        testFormatting(for: input, output, rule: .noForceUnwrapInTests)
+    }
 }