Polish UseAssertSame

Try a bit harder to find and apply the correct `JavaType.Method` so that other recipes can rely on correct type attribution.
openrewrite · Nov 20, 2024 · d16c0a7 · d16c0a7
1 parent b20da71
commit d16c0a7
Showing 1 changed file with 20 additions and 3 deletions.
diff --git a/src/main/java/org/openrewrite/java/testing/junit5/UseAssertSame.java b/src/main/java/org/openrewrite/java/testing/junit5/UseAssertSame.java
@@ -37,8 +37,8 @@ public String getDisplayName() {
 
     @Override
     public String getDescription() {
-        return "Prefers the usage of `assertSame` or `assertNotSame` methods instead of using of vanilla assertTrue " +
-               "or assertFalse with a boolean comparison.";
+        return "Prefers the usage of `assertSame` or `assertNotSame` methods instead of using of vanilla `assertTrue` " +
+               "or `assertFalse` with a boolean comparison.";
     }
 
     private static final MethodMatcher ASSERT_TRUE_MATCHER = new MethodMatcher("org.junit.jupiter.api.Assertions assertTrue(..)");
@@ -52,6 +52,8 @@ public J.MethodInvocation visitMethodInvocation(J.MethodInvocation methodInvocat
                 J.MethodInvocation mi = super.visitMethodInvocation(methodInvocation, ctx);
                 if (!ASSERT_TRUE_MATCHER.matches(mi) && !ASSERT_FALSE_MATCHER.matches(mi)) {
                     return mi;
+                } else if (mi.getMethodType() == null) {
+                    return mi;
                 }
 
                 Expression firstArgument = mi.getArguments().get(0);
@@ -73,11 +75,26 @@ public J.MethodInvocation visitMethodInvocation(J.MethodInvocation methodInvocat
                 maybeRemoveImport("org.junit.jupiter.api.Assertions");
                 maybeAddImport("org.junit.jupiter.api.Assertions", newMethodName);
 
-                JavaType.Method newType = ((JavaType.Method) mi.getName().getType()).withName(newMethodName);
+                JavaType.Method newType = assertSameMethodType(mi, newMethodName);
                 return mi.withName(mi.getName().withSimpleName(newMethodName).withType(newType))
                         .withMethodType(newType)
                         .withArguments(newArguments);
             }
+
+            private JavaType.Method assertSameMethodType(J.MethodInvocation mi, String newMethodName) {
+                JavaType.Method assertTrue = mi.getMethodType();
+                assert assertTrue != null;
+                int parameterCount = assertTrue.getParameterTypes().size();
+                JavaType.FullyQualified assertions = assertTrue.getDeclaringType();
+                for (JavaType.Method method : assertions.getMethods()) {
+                    if (method.getName().equals("assertSame") && method.getParameterNames().size() == parameterCount + 1 &&
+                        assertTrue.getParameterTypes().get(parameterCount - 1).equals(method.getParameterTypes().get(parameterCount))) {
+                        return method;
+                    }
+                }
+                // fallback when type attribution was stubbed
+                return assertTrue.withName(newMethodName);
+            }
         };
         return Preconditions.check(
                 Preconditions.or(