Spring MockMvc HAL wrapper

HalMockMvc.java

import java.util.ArrayList;
import java.util.List;
import java.util.function.Function;
import java.util.stream.Collectors;
import java.util.stream.Stream;
import org.springframework.hateoas.Link;
import org.springframework.hateoas.client.LinkDiscoverer;
import org.springframework.hateoas.mediatype.hal.forms.HalFormsLinkDiscoverer;
import org.springframework.test.web.servlet.MockMvc;
import org.springframework.test.web.servlet.ResultActions;
import org.springframework.test.web.servlet.request.MockHttpServletRequestBuilder;
import org.springframework.test.web.servlet.request.MockMvcRequestBuilders;

public class HalMockMvc {

  private final MockMvc mockMvc;
  private final String baseUri;

  public HalMockMvc(MockMvc mockMvc, String baseUri) {
    this.mockMvc = mockMvc;
    this.baseUri = baseUri;
  }

  public TraversalBuilder follow(String... relations) {
    return new TraversalBuilder(mockMvc, baseUri).follow(relations);
  }

  public TraversalBuilder follow(Hop relations) {
    return new TraversalBuilder(mockMvc, baseUri).follow(relations);
  }

  public static class TraversalBuilder {

    private final MockMvc mockMvc;
    private final String baseUri;
    private final LinkDiscoverer halLinkDiscoverer = new HalFormsLinkDiscoverer();
    private final List<Hop> hops = new ArrayList<>();

    private TraversalBuilder(MockMvc mockMvc, String baseUri) {
      this.mockMvc = mockMvc;
      this.baseUri = baseUri;
    }

    public TraversalBuilder follow(String... relation) {
      hops.addAll(Stream.of(relation).map(Hop::relation).collect(Collectors.toList()));
      return this;
    }

    public TraversalBuilder follow(Hop relation) {
      hops.add(relation);
      return this;
    }

    public ResultActions get() throws Exception {
      ResultActions lastResultActions = mockMvc.perform(MockMvcRequestBuilders.get(baseUri));
      for (Hop hop : hops) {
        String body = lastResultActions.andReturn().getResponse().getContentAsString();

        Link link = halLinkDiscoverer.findLinkWithRel(hop.relationName(), body).orElse(null);
        if (link == null) {
          throw new IllegalArgumentException(
              "Could not find relation "
                  + hop
                  + " at URI "
                  + lastResultActions.andReturn().getRequest().getRequestURI());
        }

        lastResultActions =
            mockMvc.perform(MockMvcRequestBuilders.get(link.expand(hop.parameters()).toUri()));
      }

      return lastResultActions;
    }

    public ResultActions post() throws Exception {
      return perform(MockMvcRequestBuilders::post, null);
    }

    public ResultActions post(String jsonContent) throws Exception {
      return perform(MockMvcRequestBuilders::post, jsonContent);
    }

    public ResultActions put(String jsonContent) throws Exception {
      return perform(MockMvcRequestBuilders::put, jsonContent);
    }

    public ResultActions put() throws Exception {
      return perform(MockMvcRequestBuilders::put, null);
    }

    public ResultActions delete() throws Exception {
      return perform(MockMvcRequestBuilders::delete, null);
    }

    private ResultActions perform(
        Function<String, MockHttpServletRequestBuilder> requestBuilderFactory, String jsonContent)
        throws Exception {
      ResultActions getResult = get();
      String getUri = getResult.andReturn().getRequest().getRequestURI();
      if (getUri == null) {
        throw new IllegalStateException("Expected a non null get URI");
      }
      if (jsonContent != null) {
        return mockMvc.perform(
            requestBuilderFactory
                .apply(getUri)
                .content(jsonContent)
                .contentType("application/json"));
      } else {
        return mockMvc.perform(requestBuilderFactory.apply(getUri));
      }
    }
  }
}