OpenRewrite：如何在链式/流式方法调用中替换方法？

问题分析 #

你当前的实现错误地将整个链式调用替换为isOk()，核心原因是替换坐标选择错误：链式调用在OpenRewrite的LST中是嵌套的MethodInvocation结构，isEqualTo(200)的target是前面的expectStatus()调用，而你使用method.getCoordinates().replace()会替换整个顶层MethodInvocation节点（也就是整个链式调用链），而非仅替换isEqualTo方法本身。

修正方案 #

要实现只替换链式调用末端的isEqualTo方法，需做到：

• 使用精准的替换坐标，仅替换方法名和参数，保留前面的调用链目标
• 同时处理isEqualTo(int)和isEqualTo(HttpStatusCode)两种参数场景
• 确保正确遍历嵌套调用节点

修正后的Visitor代码 #

@Override
public JavaIsoVisitor<ExecutionContext> getVisitor() {
    return new JavaIsoVisitor<ExecutionContext>() {
        // 匹配StatusAssertions下的isEqualTo方法
        private final MethodMatcher isEqualToMatcher = 
            new MethodMatcher("org.springframework.test.web.reactive.server.StatusAssertions isEqualTo(..)");
        // 匹配HttpStatusCode.valueOf(int)方法，用于处理参数为枚举的情况
        private final MethodMatcher httpStatusValueOfMatcher = 
            new MethodMatcher("org.springframework.http.HttpStatusCode valueOf(int)");
        private final JavaTemplate isOkTemplate = JavaTemplate.builder("isOk()").build();

        @Override
        public MethodInvocation visitMethodInvocation(MethodInvocation method, ExecutionContext ctx) {
            // 先处理内层调用，确保嵌套结构正确遍历
            MethodInvocation processedMethod = super.visitMethodInvocation(method, ctx);
            
            // 匹配isEqualTo方法
            if (isEqualToMatcher.matches(processedMethod.getMethodType())) {
                Expression arg = processedMethod.getArguments().get(0);
                
                // 处理 isEqualTo(200) 场景
                if (arg instanceof Literal literal 
                    && "int".equals(literal.getType().getFullyQualifiedName())
                    && 200 == (Integer) literal.getValue()) {
                    // 仅替换方法名和参数，保留前面的调用链目标
                    return isOkTemplate.apply(
                        getCursor(), 
                        processedMethod.getCoordinates().replaceMethod()
                    );
                }
                
                // 处理 isEqualTo(HttpStatusCode.valueOf(200)) 场景
                if (arg instanceof MethodInvocation statusValueCall 
                    && httpStatusValueOfMatcher.matches(statusValueCall.getMethodType())) {
                    Expression statusArg = statusValueCall.getArguments().get(0);
                    if (statusArg instanceof Literal literal 
                        && "int".equals(literal.getType().getFullyQualifiedName())
                        && 200 == (Integer) literal.getValue()) {
                        return isOkTemplate.apply(
                            getCursor(), 
                            processedMethod.getCoordinates().replaceMethod()
                        );
                    }
                }
            }
            
            return processedMethod;
        }
    };
}

关键说明 #

1. 替换坐标：使用replaceMethod()而非replace()，该坐标只会替换方法的名称和参数部分，保留方法调用的目标（即链式调用的前半部分，比如webClient.post()....expectStatus()）。
2. 嵌套遍历：先调用super.visitMethodInvocation处理内层节点，确保HttpStatusCode.valueOf(200)这类嵌套调用先被检查，再处理外层的isEqualTo。
3. 双场景处理：同时覆盖了直接传入int 200和通过HttpStatusCode.valueOf(200)传入的两种情况。

效果验证 #

对于你提供的测试代码，修正后的规则会正确生成预期结果：

import org.springframework.test.web.reactive.server.WebTestClient;

class Test {
    private final WebTestClient webClient = WebTestClient.bindToServer().build();
    void someMethod() {
      webClient
          .post()
          .uri("/some/endpoint")
          .bodyValue("someValue")
          .exchange()
          .expectStatus()
          .isOk();
    }
}

同时也能正确处理isEqualTo(HttpStatusCode.valueOf(200))的情况，将其替换为.isOk()。

内容的提问来源于stack exchange，提问作者94621