SCM-Manager/scm-test/src/main/java/sonia/scm/web/RestDispatcher.java

package sonia.scm.web;

import com.fasterxml.jackson.databind.ObjectMapper;
import org.apache.shiro.authz.AuthorizationException;
import org.apache.shiro.authz.UnauthorizedException;
import org.jboss.resteasy.mock.MockDispatcherFactory;
import org.jboss.resteasy.spi.Dispatcher;
import org.jboss.resteasy.spi.HttpRequest;
import org.jboss.resteasy.spi.HttpResponse;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import sonia.scm.AlreadyExistsException;
import sonia.scm.BadRequestException;
import sonia.scm.ConcurrentModificationException;
import sonia.scm.NotFoundException;

import javax.ws.rs.Produces;
import javax.ws.rs.core.Response;
import javax.ws.rs.core.Response.Status;
import javax.ws.rs.ext.ContextResolver;
import javax.ws.rs.ext.ExceptionMapper;
import javax.ws.rs.ext.Provider;
import java.util.HashMap;
import java.util.Map;

public class RestDispatcher {

  private static final Logger LOG = LoggerFactory.getLogger(RestDispatcher.class);

  private final Dispatcher dispatcher;
  private final EnhanceableExceptionMapper exceptionMapper;

  public RestDispatcher() {
    dispatcher = MockDispatcherFactory.createDispatcher();
    exceptionMapper = new EnhanceableExceptionMapper();
    dispatcher.getProviderFactory().register(exceptionMapper);
    dispatcher.getProviderFactory().registerProviderInstance(new JacksonProducer());
  }

  public void addSingletonResource(Object resource) {
    dispatcher.getRegistry().addSingletonResource(resource);
  }

  public void invoke(HttpRequest in, HttpResponse response) {
    dispatcher.invoke(in, response);
  }

  public void registerException(Class<? extends RuntimeException> exceptionClass, Status status) {
    exceptionMapper.registerException(exceptionClass, status);
  }

  public <T> void putDefaultContextObject(Class<T> clazz, T object) {
    dispatcher.getDefaultContextObjects().put(clazz, object);
  }

  private static class EnhanceableExceptionMapper implements ExceptionMapper<Exception> {

    private final Map<Class<? extends RuntimeException>, Integer> statusCodes = new HashMap<>();

    public EnhanceableExceptionMapper() {
      registerException(NotFoundException.class, Status.NOT_FOUND);
      registerException(AlreadyExistsException.class, Status.CONFLICT);
      registerException(ConcurrentModificationException.class, Status.CONFLICT);
      registerException(UnauthorizedException.class, Status.FORBIDDEN);
      registerException(AuthorizationException.class, Status.FORBIDDEN);
      registerException(BadRequestException.class, Status.BAD_REQUEST);
    }

    private void registerException(Class<? extends RuntimeException> exceptionClass, Status status) {
      statusCodes.put(exceptionClass, status.getStatusCode());
    }

    @Override
    public Response toResponse(Exception e) {
      return Response.status(getStatus(e)).entity(e.getMessage()).build();
    }

    private Integer getStatus(Exception ex) {
      return statusCodes
        .entrySet()
        .stream()
        .filter(e -> e.getKey().isAssignableFrom(ex.getClass()))
        .map(Map.Entry::getValue)
        .findAny()
        .orElse(handleUnknownException(ex));
    }

    private Integer handleUnknownException(Exception ex) {
      LOG.info("got unknown exception in rest api test", ex);
      return 500;
    }
  }

  @Provider
  @Produces("application/*+json")
  public static class JacksonProducer implements ContextResolver<ObjectMapper> {
    public JacksonProducer() {
      this.json
        = new ObjectMapper().findAndRegisterModules();
    }

    @Override
    public ObjectMapper getContext(Class<?> objectType) {
      return json;
    }

    private final ObjectMapper json;
  }
}