Slight polish

openrewrite · Sep 19, 2024 · 90339bb · 90339bb
1 parent 6902791
commit 90339bb
Show file tree

Hide file tree

Showing 2 changed files with 29 additions and 46 deletions.
diff --git a/src/main/java/org/openrewrite/staticanalysis/RemoveMethodCallVisitor.java b/src/main/java/org/openrewrite/staticanalysis/RemoveMethodCallVisitor.java
@@ -22,7 +22,6 @@
 import org.openrewrite.java.tree.*;
 
 import java.util.function.BiPredicate;
-import java.util.function.Supplier;
 
 /**
  * Removes all {@link MethodCall} matching both the
@@ -69,8 +68,8 @@ public class RemoveMethodCallVisitor<P> extends JavaIsoVisitor<P> {
             }
 
             // If the method invocation is in a fluent chain, remove just the current invocation
-            if (isInFluentChain(method)) {
-                //noinspection ConstantConditions
+            if (method.getSelect() instanceof J.MethodInvocation &&
+                TypeUtils.isOfType(method.getType(), method.getSelect().getType())) {
                 return super.visitMethodInvocation((J.MethodInvocation) method.getSelect(), p);
             }
         }
@@ -94,19 +93,4 @@ private boolean isStatementInParentBlock(Statement method) {
         J.Block parentBlock = getCursor().firstEnclosing(J.Block.class);
         return parentBlock == null || parentBlock.getStatements().contains(method);
     }
-
-    private boolean isInFluentChain(J.MethodInvocation method) {
-        Expression select = method.getSelect();
-        if (select instanceof J.MethodInvocation) {
-            return sameReturnType(method, (J.MethodInvocation) select);
-        }
-        return false;
-    }
-
-    private boolean sameReturnType(J.MethodInvocation m1, J.MethodInvocation m2) {
-        if (m1.getMethodType() == null || m2.getMethodType() == null) {
-            return false;
-        }
-        return m1.getMethodType().getReturnType() == m2.getMethodType().getReturnType();
-    }
 }
diff --git a/src/test/java/org/openrewrite/staticanalysis/RemoveMethodCallVisitorTest.java b/src/test/java/org/openrewrite/staticanalysis/RemoveMethodCallVisitorTest.java
@@ -17,7 +17,6 @@
 
 import org.junit.jupiter.api.Test;
 import org.openrewrite.DocumentExample;
-import org.openrewrite.ExecutionContext;
 import org.openrewrite.java.MethodMatcher;
 import org.openrewrite.java.tree.J;
 import org.openrewrite.test.RecipeSpec;
@@ -44,7 +43,7 @@ void assertTrueIsRemoved() {
             """
               abstract class Test {
                   abstract void assertTrue(boolean condition);
-
+              
                   void test() {
                       System.out.println("Hello");
                       assertTrue(true);
@@ -54,7 +53,7 @@ void test() {
               """, """
               abstract class Test {
                   abstract void assertTrue(boolean condition);
-
+              
                   void test() {
                       System.out.println("Hello");
                       System.out.println("World");
@@ -73,7 +72,7 @@ void assertTrueFalseIsNotRemoved() {
             """
               abstract class Test {
                   abstract void assertTrue(boolean condition);
-
+              
                   void test() {
                       System.out.println("Hello");
                       assertTrue(false);
@@ -96,7 +95,7 @@ void assertTrueTwoArgIsRemoved() {
             """
               abstract class Test {
                   abstract void assertTrue(String message, boolean condition);
-
+              
                   void test() {
                       System.out.println("Hello");
                       assertTrue("message", true);
@@ -107,7 +106,7 @@ void test() {
             """
               abstract class Test {
                   abstract void assertTrue(String message, boolean condition);
-
+              
                   void test() {
                       System.out.println("Hello");
                       System.out.println("World");
@@ -126,7 +125,7 @@ void doesNotRemoveAssertTrueIfReturnValueIsUsed() {
             """
               abstract class Test {
                   abstract int assertTrue(boolean condition);
-
+              
                   void test() {
                       System.out.println("Hello");
                       int value = assertTrue(true);
@@ -140,31 +139,31 @@ void test() {
 
     @Test
     void removeMethodCallFromFluentChain() {
-        final MethodMatcher methodMatcher = new MethodMatcher("java.lang.StringBuilder append(..)");
-        final RemoveMethodCallVisitor<ExecutionContext> visitor =
-          new RemoveMethodCallVisitor<>(methodMatcher, (i, e) -> true);
         rewriteRun(
-          spec -> spec.recipe(RewriteTest.toRecipe(() -> visitor)),
+          spec -> spec.recipe(RewriteTest.toRecipe(() -> new RemoveMethodCallVisitor<>(
+            new MethodMatcher("java.lang.StringBuilder append(..)"), (i, e) -> true))),
           // language=java
           java(
             """
-            class Main {
-                void hello() {
-                    final String s = new StringBuilder("hello")
-                            .delete(1, 2)
-                            .append("world")
-                            .toString();
-                }
-            }
-            """,
+              class Main {
+                  void hello() {
+                      final String s = new StringBuilder("hello")
+                              .delete(1, 2)
+                              .append("world")
+                              .toString();
+                  }
+              }
+              """,
             """
-            class Main {
-                void hello() {
-                    final String s = new StringBuilder("hello")
-                            .delete(1, 2)
-                            .toString();
-                }
-            }
-            """));
+              class Main {
+                  void hello() {
+                      final String s = new StringBuilder("hello")
+                              .delete(1, 2)
+                              .toString();
+                  }
+              }
+              """
+          )
+        );
     }
 }