From 8c272b9e634427714b9177daf12fc0e95254a678 Mon Sep 17 00:00:00 2001 From: Tim te Beek Date: Tue, 25 Aug 2026 14:23:04 +0200 Subject: [PATCH] Retain the original stubbing method when rewriting static `when` to `MockedStatic` `MockitoWhenOnStaticToMockStatic` matched any `when(Static.method()).xxx(...)` statement but always emitted `.thenReturn(firstArgument)`. That silently turned `thenAnswer`/`thenThrow`/`then`/`thenCallRealMethod` into `thenReturn`, and dropped all but the first argument of `thenReturn(a, b, c)`. Carry the invoked method name and its full argument list into the Java and Kotlin templates instead. Fixes #1102 --- .../MockitoWhenOnStaticToMockStatic.java | 42 ++- .../MockitoWhenOnStaticToMockStaticTest.java | 256 ++++++++++++++++++ 2 files changed, 285 insertions(+), 13 deletions(-) diff --git a/src/main/java/org/openrewrite/java/testing/mockito/MockitoWhenOnStaticToMockStatic.java b/src/main/java/org/openrewrite/java/testing/mockito/MockitoWhenOnStaticToMockStatic.java index 53bc41e20..4965a84fd 100644 --- a/src/main/java/org/openrewrite/java/testing/mockito/MockitoWhenOnStaticToMockStatic.java +++ b/src/main/java/org/openrewrite/java/testing/mockito/MockitoWhenOnStaticToMockStatic.java @@ -31,7 +31,9 @@ import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicReference; +import static java.util.Arrays.asList; import static java.util.Collections.emptyList; +import static java.util.Collections.singletonList; import static java.util.Objects.requireNonNull; import static org.openrewrite.java.VariableNameUtils.GenerationStrategy.INCREMENT_NUMBER; import static org.openrewrite.java.VariableNameUtils.generateVariableName; @@ -150,13 +152,13 @@ private J.Try tryWithMockedStatic(J.Block block, List statements, Int Map pendingResources) { String className = invokedType.getClassName(); String variableName = generateVariableName("mock" + className + ++varCounter, updateCursor(block), INCREMENT_NUMBER); - Expression thenReturnArg = statement.getArguments().get(0); + List stubbingArguments = stubbingArguments(statement); J.Try try_ = (J.Try) javaTemplateMockStatic(String.format( "try(MockedStatic<%1$s> %2$s = mockStatic(%1$s.class)) {\n" + - " %2$s.when(() -> #{any()}).thenReturn(#{any()});\n" + - "}", className, variableName), ctx) - .apply(getCursor(), block.getCoordinates().firstStatement(), whenArg, thenReturnArg) + " %2$s.when(() -> #{any()}).%3$s(%4$s);\n" + + "}", className, variableName, statement.getSimpleName(), argumentPlaceholders(stubbingArguments)), ctx) + .apply(getCursor(), block.getCoordinates().firstStatement(), templateParameters(singletonList(whenArg), stubbingArguments)) .getStatements().get(0); List precedingStatements = statements.subList(0, index); @@ -187,8 +189,10 @@ private Object nameForReuse(J.Block block, String className, J.VariableDeclarati private Statement reuseMockedStatic(J.Block block, J.MethodInvocation statement, Object variable, J.MethodInvocation whenArg, ExecutionContext ctx) { String mockedStaticVariableTemplate = variable instanceof J ? "#{any()}" : "#{}"; J.Block cursorBlock = (J.Block) getCursor().getValue(); - Statement replacement = javaTemplateMockStatic(mockedStaticVariableTemplate + ".when(() -> #{any()}).thenReturn(#{any()});", ctx) - .apply(getCursor(), cursorBlock.getCoordinates().firstStatement(), variable, whenArg, statement.getArguments().get(0)) + List stubbingArguments = stubbingArguments(statement); + Statement replacement = javaTemplateMockStatic(String.format("%s.when(() -> #{any()}).%s(%s);", + mockedStaticVariableTemplate, statement.getSimpleName(), argumentPlaceholders(stubbingArguments)), ctx) + .apply(getCursor(), cursorBlock.getCoordinates().firstStatement(), templateParameters(asList(variable, whenArg), stubbingArguments)) .getStatements().get(0); return replacement.withPrefix(statement.getPrefix()); } @@ -200,12 +204,12 @@ private List mockedStatic(J.Block block, J.MethodInvocation statement // We know it will have a matching `@Before*` annotation based on callers String matchedAnnotation = requireNonNull(tryGetMatchedAnnotationOnMethodDeclaration(containingMethod, BEFORE)); String correspondingAfterFqn = matchedAnnotation.replace(".Before", ".After"); - Expression thenReturnArg = statement.getArguments().get(0); + List stubbingArguments = stubbingArguments(statement); List statements = javaTemplateMockStatic(String.format( "%2$s = mockStatic(%1$s.class);\n" + - "%2$s.when(() -> #{any()}).thenReturn(#{any()});", className, variableName), ctx) - .apply(getCursor(), block.getCoordinates().firstStatement(), whenArg, thenReturnArg) + "%2$s.when(() -> #{any()}).%3$s(%4$s);", className, variableName, statement.getSimpleName(), argumentPlaceholders(stubbingArguments)), ctx) + .apply(getCursor(), block.getCoordinates().firstStatement(), templateParameters(singletonList(whenArg), stubbingArguments)) .getStatements().subList(0, 2); doAfterVisit(new JavaIsoVisitor() { @@ -307,14 +311,14 @@ public J.MethodInvocation visitMethodInvocation(J.MethodInvocation method, Execu } String returnTypeName = getReturnTypeName(whenArg); - Expression thenReturnArg = m.getArguments().get(0); + List stubbingArguments = stubbingArguments(m); J.MethodInvocation rewritten = KotlinTemplate.builder(String.format( - "%s.`when`<%s> { #{any()} }.thenReturn(#{any()})", - paramName, returnTypeName)) + "%s.`when`<%s> { #{any()} }.%s(%s)", + paramName, returnTypeName, m.getSimpleName(), argumentPlaceholders(stubbingArguments))) .imports("org.mockito.MockedStatic") .parser(KotlinParser.builder().classpathFromResources(ctx, "mockito-core-5")) .build() - .apply(getCursor(), m.getCoordinates().replace(), whenArg, thenReturnArg); + .apply(getCursor(), m.getCoordinates().replace(), templateParameters(singletonList(whenArg), stubbingArguments)); maybeRemoveImport("org.mockito.Mockito.when"); return rewritten; } @@ -334,6 +338,18 @@ public J.MethodInvocation visitMethodInvocation(J.MethodInvocation method, Execu return null; } + private static List stubbingArguments(J.MethodInvocation statement) { + return ListUtils.filter(statement.getArguments(), argument -> !(argument instanceof J.Empty)); + } + + private static String argumentPlaceholders(List arguments) { + return String.join(", ", Collections.nCopies(arguments.size(), "#{any()}")); + } + + private static Object[] templateParameters(List leading, List stubbingArguments) { + return ListUtils.concatAll(leading, stubbingArguments).toArray(); + } + private static JavaType.@Nullable Class getTypeFromInvocation(J.MethodInvocation whenArg) { J.Identifier clazz = null; // Having a fieldType implies that something is a field rather than a class itself diff --git a/src/test/java/org/openrewrite/java/testing/mockito/MockitoWhenOnStaticToMockStaticTest.java b/src/test/java/org/openrewrite/java/testing/mockito/MockitoWhenOnStaticToMockStaticTest.java index 78f3a3aab..033e22225 100644 --- a/src/test/java/org/openrewrite/java/testing/mockito/MockitoWhenOnStaticToMockStaticTest.java +++ b/src/test/java/org/openrewrite/java/testing/mockito/MockitoWhenOnStaticToMockStaticTest.java @@ -457,6 +457,178 @@ void tearDown() { ); } + @Test + void retainsThenAnswer() { + rewriteRun( + //language=java + java( + """ + import org.example.A; + + import static org.mockito.Mockito.*; + + class Test { + void test() { + when(A.getNumber()).thenAnswer(invocation -> -1); + } + } + """, + """ + import org.example.A; + import org.mockito.MockedStatic; + + import static org.mockito.Mockito.*; + + class Test { + void test() { + try (MockedStatic mockA1 = mockStatic(A.class)) { + mockA1.when(() -> A.getNumber()).thenAnswer(invocation -> -1); + } + } + } + """ + ) + ); + } + + @Test + void retainsThenThrow() { + rewriteRun( + //language=java + java( + """ + import org.example.A; + + import static org.mockito.Mockito.*; + + class Test { + void test() { + when(A.getNumber()).thenThrow(new IllegalStateException()); + } + } + """, + """ + import org.example.A; + import org.mockito.MockedStatic; + + import static org.mockito.Mockito.*; + + class Test { + void test() { + try (MockedStatic mockA1 = mockStatic(A.class)) { + mockA1.when(() -> A.getNumber()).thenThrow(new IllegalStateException()); + } + } + } + """ + ) + ); + } + + @Test + void retainsAllConsecutiveReturnValues() { + rewriteRun( + //language=java + java( + """ + import org.example.A; + + import static org.mockito.Mockito.*; + + class Test { + void test() { + when(A.getNumber()).thenReturn(-1, -2, -3); + } + } + """, + """ + import org.example.A; + import org.mockito.MockedStatic; + + import static org.mockito.Mockito.*; + + class Test { + void test() { + try (MockedStatic mockA1 = mockStatic(A.class)) { + mockA1.when(() -> A.getNumber()).thenReturn(-1, -2, -3); + } + } + } + """ + ) + ); + } + + @Test + void retainsThenCallRealMethodWithoutArguments() { + rewriteRun( + //language=java + java( + """ + import org.example.A; + + import static org.mockito.Mockito.*; + + class Test { + void test() { + when(A.getNumber()).thenCallRealMethod(); + } + } + """, + """ + import org.example.A; + import org.mockito.MockedStatic; + + import static org.mockito.Mockito.*; + + class Test { + void test() { + try (MockedStatic mockA1 = mockStatic(A.class)) { + mockA1.when(() -> A.getNumber()).thenCallRealMethod(); + } + } + } + """ + ) + ); + } + + @Test + void retainsStubbingMethodWhenReusingMockedStatic() { + rewriteRun( + //language=java + java( + """ + import org.example.A; + + import static org.mockito.Mockito.*; + + class Test { + void test() { + when(A.getNumber()).thenReturn(-1); + when(A.getNumber()).thenThrow(new IllegalStateException()); + } + } + """, + """ + import org.example.A; + import org.mockito.MockedStatic; + + import static org.mockito.Mockito.*; + + class Test { + void test() { + try (MockedStatic mockA1 = mockStatic(A.class)) { + mockA1.when(() -> A.getNumber()).thenReturn(-1); + mockA1.when(() -> A.getNumber()).thenThrow(new IllegalStateException()); + } + } + } + """ + ) + ); + } + @Nested class UsingJunit4 { @Test @@ -962,6 +1134,51 @@ void test1() { ); } + @Test + void retainsThenThrow_inBeforeEach() { + rewriteRun( + //language=java + java( + """ + import org.example.A; + import org.junit.jupiter.api.BeforeEach; + + import static org.mockito.Mockito.*; + + class Test { + @BeforeEach + public void setUp() { + when(A.getNumber()).thenThrow(new IllegalStateException()); + } + } + """, + """ + import org.example.A; + import org.junit.jupiter.api.AfterEach; + import org.junit.jupiter.api.BeforeEach; + import org.mockito.MockedStatic; + + import static org.mockito.Mockito.*; + + class Test { + private MockedStatic mockA1; + + @BeforeEach + public void setUp() { + mockA1 = mockStatic(A.class); + mockA1.when(() -> A.getNumber()).thenThrow(new IllegalStateException()); + } + + @AfterEach + public void tearDown() { + mockA1.close(); + } + } + """ + ) + ); + } + @Test void handlesStaticMocks_inBeforeEach_withExistingAfterEach() { rewriteRun( @@ -2166,6 +2383,45 @@ fun testStaticMethod() { ); } + @Test + void shouldRetainKotlinStubbingMethodOtherThanThenReturn() { + rewriteRun( + spec -> spec.afterTypeValidationOptions(TypeValidation.none()), + //language=kotlin + kotlin( + """ + import org.junit.jupiter.api.Test + import org.mockito.Mockito.`when` + import org.mockito.Mockito.mockStatic + import java.util.Calendar + + class MyTest { + @Test + fun testStaticMethod() { + mockStatic(Calendar::class.java).use { + `when`(Calendar.getInstance()).thenThrow(IllegalStateException()) + } + } + } + """, + """ + import org.junit.jupiter.api.Test + import org.mockito.Mockito.mockStatic + import java.util.Calendar + + class MyTest { + @Test + fun testStaticMethod() { + mockStatic(Calendar::class.java).use { + it.`when` { Calendar.getInstance() }.thenThrow(IllegalStateException()) + } + } + } + """ + ) + ); + } + @Test void shouldNotRewriteKotlinMockitoWhenWhenStaticClassDoesNotMatchUseReceiver() { rewriteRun(