[FIXED] Mock flatMap ServerResponse Body in Spring Webflux

Issue

I am running a Spring Webflux application, as such I am using flatMap to convert my request body. I have a Router that looks like this,

@Configuration
public class TallyRouter {
    private static final String TALLY_BASE_URL = "/admin/tally";
    private static final String TALLY_ID_URL = "/{tallyId}";

    @Bean
    RouterFunction<ServerResponse> tallyRoutes(tallyHandler handler) {
        return RouterFunctions.route(GET(TALLY_CONFIG_BASE_URL + "/active"), handler::getActiveTally)
                .andRoute(POST(TALLY_BASE_URL).and(accept(MediaType.APPLICATION_JSON)), handler::createTally);
}

My handler class looks like this,

@Component
@RequiredArgsConstructor
public class TallyHandler {

    @Autowired
    TallyService tallyService;

    private static final Logger logger = LogManager.getLogger(TallyHandler.class);

    public Mono<ServerResponse> createTally(ServerRequest request) {
        Mono<NewTallyRequest> req = request.bodyToMono(NewTallyRequest.class);
        return req
                .doOnNext(result -> logger.info(result))
                .flatMap(a -> ServerResponse.ok().contentType(MediaType.APPLICATION_JSON)
                        .body(tallyService.createTally(a, request.exchange()), String.class));   // the line of the code that is null
    }
}

This is how my test looks like,

@ExtendWith(MockitoExtension.class)
@WebFluxTest
@ContextConfiguration(classes = { TallyHandler.class, TallyRouter.class })
@ActiveProfiles("test")
class TallyHandlerTest {

    @MockBean
    WebClient webClient;

    @MockBean
    TallyService tallyService;

    @Autowired
    private ApplicationContext context;

    WebTestClient webTestClient;

    MockWebServer mockWebServer;

    String res = "testRes";

    @BeforeEach
    void before() throws IOException {

        webTestClient = WebTestClient.bindToApplicationContext(this.context).configureClient().build();

        mockWebServer = new MockWebServer();
        mockWebServer.start();
    }

    @AfterEach
    void tearDown() throws IOException {
        mockWebServer.close();
    }

    @Test
    @WithMockUser(roles = "ADMIN")
    void createTallyTest_200() {
        String url = "/admin/tally";

        MockResponse mockedResponse = new MockResponse().setStatus("200").setBody(res);
        mockWebServer.enqueue(mockedResponse);

        NewTallyRequest req = new NewTallyRequest();
        req.setIsActive(isActive);

        when(tallyService.createTally(req, exchange)).thenReturn(Mono.just(res));

        webTestClient.mutateWith(csrf()).post().uri(uriBuilder -> uriBuilder.path(url).build())
                .body(BodyInserters.fromValue(req)).exchange().expectStatus().is2xxSuccessful();
    }
}

Error:

java.lang.IllegalArgumentException: 'publisher' must not be null
    at org.springframework.util.Assert.notNull(Assert.java:201)
    Suppressed: reactor.core.publisher.FluxOnAssembly$OnAssemblyException: 
Error has been observed at the following site(s):
    *__checkpoint ⇢ org.springframework.security.web.server.authorization.AuthorizationWebFilter [DefaultWebFilterChain]
    *__checkpoint ⇢ org.springframework.security.web.server.authorization.ExceptionTranslationWebFilter [DefaultWebFilterChain]
    *__checkpoint ⇢ org.springframework.security.web.server.authentication.logout.LogoutWebFilter [DefaultWebFilterChain]
    *__checkpoint ⇢ org.springframework.security.web.server.savedrequest.ServerRequestCacheWebFilter [DefaultWebFilterChain]
    *__checkpoint ⇢ org.springframework.security.web.server.context.SecurityContextServerWebExchangeWebFilter [DefaultWebFilterChain]
    *__checkpoint ⇢ org.springframework.security.web.server.ui.LogoutPageGeneratingWebFilter [DefaultWebFilterChain]
    *__checkpoint ⇢ org.springframework.security.web.server.ui.LoginPageGeneratingWebFilter [DefaultWebFilterChain]
    *__checkpoint ⇢ org.springframework.security.web.server.authentication.AuthenticationWebFilter [DefaultWebFilterChain]
    *__checkpoint ⇢ org.springframework.security.web.server.authentication.AuthenticationWebFilter [DefaultWebFilterChain]
    *__checkpoint ⇢ org.springframework.security.web.server.context.ReactorContextWebFilter [DefaultWebFilterChain]
    *__checkpoint ⇢ org.springframework.security.web.server.csrf.CsrfWebFilter [DefaultWebFilterChain]
    *__checkpoint ⇢ org.springframework.security.web.server.header.HttpHeaderWriterWebFilter [DefaultWebFilterChain]
    *__checkpoint ⇢ org.springframework.security.config.web.server.ServerHttpSecurity$ServerWebExchangeReactorContextWebFilter [DefaultWebFilterChain]
    *__checkpoint ⇢ org.springframework.security.web.server.WebFilterChainProxy [DefaultWebFilterChain]
    *__checkpoint ⇢ org.springframework.security.web.server.csrf.CsrfWebFilter [DefaultWebFilterChain]
    *__checkpoint ⇢ HTTP POST "/admin/tally" [ExceptionHandlingWebHandler]
Original Stack Trace:
        at org.springframework.util.Assert.notNull(Assert.java:201)
        at org.springframework.web.reactive.function.BodyInserters.fromPublisher(BodyInserters.java:182)
        at org.springframework.web.reactive.function.server.DefaultServerResponseBuilder.body(DefaultServerResponseBuilder.java:233)
        at my.company.app.handler.TallyHandler.lambda$1(TallyHandler.java:58)

I have tried to mock the return value of the tallyService but the test is not picking it up in the body of ServerResponse. What should I do to mock that part of the code to return a correct ServerResponse value?

Solution

I finally found the answer to this problem, I didn't mock the body properly.

I was doing this,

when(tallyService.createTally(req, exchange)).thenReturn(Mono.just(res));

when I should be doing this,

when(tallyService.createTally(any(NewTallyRequest.class), any(ServerWebExchange.class)).thenReturn(Mono.just(res));

The body of the ServerResponse entity wasn't mocked hence I was faced with the error of returning an empty body.

Answered By - Seng Wee
Answer Checked By - David Marino (JavaFixing Volunteer)