OpenRewrite – How to replace a method in a chained/fluent method invocation?

I am currently getting started with writing rewrite recipes and am struggling a bit with replacing methods in a fluent interface. In particular, I’d like to replace .isEqualTo(200) (and once that works also .isEqualTo(HttpStatusCode.valueOf(200)) with .isOk():

Before:

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()
          .isEqualTo(200);
    }
}

After:

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();
    }
}

I have come up with the following visitor that is matching the isEqualTo method and replacing it with isOk():

@Override
public JavaIsoVisitor<ExecutionContext> getVisitor() {
    return new JavaIsoVisitor<ExecutionContext>() {

        private final MethodMatcher methodMatcher
            = new MethodMatcher("org.springframework.test.web.reactive.server.StatusAssertions isEqualTo(..)");
        private final JavaType JAVA_TYPE_INT = JavaType.buildType("int");
        private final JavaTemplate isOkTemplate 
            = JavaTemplate.builder("isOk()").build();

        @Override
        public MethodInvocation visitMethodInvocation(MethodInvocation method, ExecutionContext p) {
            // exit if method doesn't match isEqualTo(..)
            if (!methodMatcher.matches(method.getMethodType())) {
                return method;
            }
            Expression expression = method.getArguments().get(0);
            // isEqualTo has two signatures: isEqualTo(int) and isEqualTo(HttpStatusCode)
            // handle `isEqualTo(int) here
            if(expression instanceof Literal) {
                if(JAVA_TYPE_INT.equals(expression.getType())) {
                    Literal literal = (Literal) expression;
                    if(literal.getValue() instanceof Integer) {
                        if ((int) literal.getValue() == 200) {
                            MethodInvocation m = isOkTemplate.apply(getCursor(), method.getCoordinates().replace());
                            return m;
                        }
                    }
                    return method;
                }
                return method;
            }
            return super.visitMethodInvocation(method, p);
        }
    };
}

This implementation is replacing the whole webClient invocation (webClient.post().(..).isOk()) with isOk():

import org.springframework.test.web.reactive.server.WebTestClient;
  
class Test {
    private final WebTestClient webClient = WebTestClient.bindToServer().build();
    void someMethod() {
        isOk();
    }
}

Having the MethodInvocation nesting in mind that is described in the LST examples document, I would guess that my MethodMatcher is too specific and I would actually need to capture the outmost invocation and somehow stream() over it and replace the nested isEqualTo invocation.
Can someone point me in the right direction?

>Solution :

I’d say you’re really close; What you’re currently doing is replacing the whole method, and with that the select that it’s called upon. Did you already try replaceMethod instead of replace for the coordinates?

Alternatively you do something like return m.withSelect(method.getSelect()) to retain the chain before.