Support member references in AssertJ assertThrows (#330)

Author: Meijuh

src/main/java/org/openrewrite/java/testing/assertj/JUnitAssertThrowsToAssertExceptionType.java

@@ -67,23 +67,34 @@ private Supplier<JavaParser> assertionsParser(ExecutionContext ctx) {
         }
         private static final MethodMatcher ASSERT_THROWS_MATCHER = new MethodMatcher("org.junit.jupiter.api.Assertions assertThrows(..)");
 
+        private static final JavaType THROWING_CALLABLE_TYPE = JavaType.buildType("org.assertj.core.api.ThrowableAssert.ThrowingCallable");
+
         @Override
         public J.MethodInvocation visitMethodInvocation(J.MethodInvocation method, ExecutionContext ctx) {
             J.MethodInvocation mi = super.visitMethodInvocation(method, ctx);
             if (ASSERT_THROWS_MATCHER.matches(mi) && mi.getArguments().size() == 2) {
-                J.Lambda lambdaArg = (J.Lambda) mi.getArguments().get(1);
-                lambdaArg = lambdaArg.withType(JavaType.buildType("org.assertj.core.api.ThrowableAssert.ThrowingCallable"));
-                mi = mi.withTemplate(
-                        JavaTemplate
-                                .builder(this::getCursor,
-                                        "assertThatExceptionOfType(#{any(java.lang.Class)}).isThrownBy(#{any(org.assertj.core.api.ThrowableAssert.ThrowingCallable)})")
-                                .javaParser(assertionsParser(ctx))
-                                .staticImports("org.assertj.core.api.AssertionsForClassTypes.assertThatExceptionOfType")
-                                .build(),
-                        mi.getCoordinates().replace(),
-                        mi.getArguments().get(0), lambdaArg);
-                maybeAddImport("org.assertj.core.api.AssertionsForClassTypes", "assertThatExceptionOfType");
-                maybeRemoveImport("org.junit.jupiter.api.Assertions.assertThrows");
+                J executable = mi.getArguments().get(1);
+                if (executable instanceof J.Lambda) {
+                    executable = ((J.Lambda) executable).withType(THROWING_CALLABLE_TYPE);
+                } else if (executable instanceof J.MemberReference) {
+                    executable = ((J.MemberReference) executable).withType(THROWING_CALLABLE_TYPE);
+                } else {
+                    executable = null;
+                }
+
+                if (executable != null) {
+                    mi = mi.withTemplate(
+                      JavaTemplate
+                        .builder(this::getCursor,
+                          "assertThatExceptionOfType(#{any(java.lang.Class)}).isThrownBy(#{any(org.assertj.core.api.ThrowableAssert.ThrowingCallable)})")
+                        .javaParser(assertionsParser(ctx))
+                        .staticImports("org.assertj.core.api.AssertionsForClassTypes.assertThatExceptionOfType")
+                        .build(),
+                      mi.getCoordinates().replace(),
+                      mi.getArguments().get(0), executable);
+                    maybeAddImport("org.assertj.core.api.AssertionsForClassTypes", "assertThatExceptionOfType");
+                    maybeRemoveImport("org.junit.jupiter.api.Assertions.assertThrows");
+                }
             }
             return mi;
         }

src/test/java/org/openrewrite/java/testing/assertj/JUnitAssertThrowsToAssertExceptionTypeTest.java

@@ -63,4 +63,39 @@ public void throwsExceptionWithSpecificType() {
           )
         );
     }
+
+    @Test
+    void memberReference() {
+        //language=java
+        rewriteRun(
+          java(
+            """
+              import static org.junit.jupiter.api.Assertions.assertThrows;
+              import java.util.concurrent.CompletableFuture;
+              import java.util.concurrent.ExecutionException;
+
+              public class MemberReferenceTest {
+
+                  public void throwsWithMemberReference() {
+                      CompletableFuture<Boolean> future = new CompletableFuture<>();
+                      assertThrows(ExecutionException.class, future::get);
+                  }
+              }
+              """,
+            """
+              import static org.assertj.core.api.AssertionsForClassTypes.assertThatExceptionOfType;
+              import java.util.concurrent.CompletableFuture;
+              import java.util.concurrent.ExecutionException;
+
+              public class MemberReferenceTest {
+
+                  public void throwsWithMemberReference() {
+                      CompletableFuture<Boolean> future = new CompletableFuture<>();
+                      assertThatExceptionOfType(ExecutionException.class).isThrownBy(future::get);
+                  }
+              }
+              """
+          )
+        );
+    }
 }