JasonLeyba: Allow the UrlMapper to configure mime-type specific globa…

…l handlers.

r14293

java/server/src/org/openqa/selenium/remote/server/DriverServlet.java

@@ -117,6 +117,7 @@
 import org.openqa.selenium.remote.server.renderer.JsonResult;
 import org.openqa.selenium.remote.server.renderer.RedirectResult;
 import org.openqa.selenium.remote.server.rest.Handler;
+import org.openqa.selenium.remote.server.rest.Result;
 import org.openqa.selenium.remote.server.rest.ResultConfig;
 import org.openqa.selenium.remote.server.rest.ResultType;
 import org.openqa.selenium.remote.server.rest.UrlMapper;
@@ -183,12 +184,11 @@ private void setupMappings(DriverSessions driverSessions, Logger logger) {
     postMapper = new UrlMapper(driverSessions, logger);
     deleteMapper = new UrlMapper(driverSessions, logger);
 
-    getMapper.addGlobalHandler(ResultType.EXCEPTION,
-        new JsonErrorExceptionResult(EXCEPTION, RESPONSE));
-    postMapper.addGlobalHandler(ResultType.EXCEPTION,
-        new JsonErrorExceptionResult(EXCEPTION, RESPONSE));
-    deleteMapper.addGlobalHandler(ResultType.EXCEPTION,
+    final Result jsonErrorResult = new Result("",
         new JsonErrorExceptionResult(EXCEPTION, RESPONSE));
+    getMapper.addGlobalHandler(ResultType.EXCEPTION, jsonErrorResult);
+    postMapper.addGlobalHandler(ResultType.EXCEPTION, jsonErrorResult);
+    deleteMapper.addGlobalHandler(ResultType.EXCEPTION, jsonErrorResult);
 
     postMapper.bind("/config/drivers", AddConfig.class)
         .on(ResultType.SUCCESS, emptyResponse);

java/server/src/org/openqa/selenium/remote/server/rest/ResultConfig.java

@@ -20,7 +20,10 @@
 import static com.google.common.base.Preconditions.checkNotNull;
 
 import com.google.common.annotations.VisibleForTesting;
+import com.google.common.collect.LinkedHashMultimap;
 import com.google.common.collect.Lists;
+import com.google.common.collect.Multimap;
+import com.google.common.collect.Multimaps;
 
 import org.openqa.selenium.WebDriverException;
 import org.openqa.selenium.remote.ErrorCodes;
@@ -41,6 +44,7 @@
 import java.lang.reflect.Constructor;
 import java.lang.reflect.Method;
 import java.lang.reflect.UndeclaredThrowableException;
+import java.util.Collection;
 import java.util.HashMap;
 import java.util.Iterator;
 import java.util.LinkedHashSet;
@@ -60,8 +64,7 @@ public class ResultConfig {
   private final String[] sections;
   private final HandlerFactory handlerFactory;
   private final DriverSessions sessions;
-  private final Map<ResultType, Set<Result>> resultToRender =
-      new HashMap<ResultType, Set<Result>>();
+  private final Multimap<ResultType, Result> resultToRender = LinkedHashMultimap.create();
   private final String url;
   private final Logger log;
 
@@ -124,28 +127,47 @@ protected Handler populate(Handler handler, String pathString) {
     return handler;
   }
 
-  public ResultConfig on(ResultType success, Renderer renderer) {
-    return on(success, renderer, "");
+  /**
+   * Configures this instance to handle a particular type of result with the given renderer. This
+   * result handler will be registered with an empty mime-type.  Accordingly, it will only be used
+   * if there are no other handlers registered with an exact mime-type match.
+   *
+   * @param resultType The type of result to configure.
+   * @param renderer The renderer to use.
+   * @return A self reference for fluency.
+   * @see #on(ResultType, Result)
+   */
+  public ResultConfig on(ResultType resultType, Renderer renderer) {
+    return on(resultType, renderer, "");
   }
 
   /*
    * Configure this ResultConfig to handle results of type ResultType with a specific renderer. The
    * mimeType is used to distinguish between JSON calls and "ordinary" browser pointed at the remote
    * WD Server, which is not implemented at all yet.
+   * @see #on(ResultType, Result)
    */
   public ResultConfig on(ResultType success, Renderer renderer, String mimeType) {
-    Set<Result> results = resultToRender.get(success);
-    if (results == null) {
-      results = new LinkedHashSet<Result>();
-      resultToRender.put(success, results);
-    }
+    return on(success, new Result(mimeType, renderer));
+  }
 
+  /**
+   * Configures how this instance will handle specific types of results. Each ResultType may be
+   * handled by multiple Results. Upon rendering a response, this instance will select the first
+   * Result that is an exact mime-type match for the original HTTP request (results are checked in
+   * the order registered). There may only be one Result registered for each mime-type.
+   *
+   * @param type The type of result to configure for.
+   * @param result The handler for the given result type.
+   * @return A self reference for fluency.
+   */
+  public ResultConfig on(ResultType type, Result result) {
     // There should not be more than one renderer for each result and
     // mime type.
-    for (Result existingResult : results) {
-      assert(!existingResult.isExactMimeTypeMatch(mimeType));
+    for (Result existingResult : resultToRender.get(type)) {
+      assert(!existingResult.isExactMimeTypeMatch(result.getMimeType()));
     }
-    results.add(new Result(mimeType, renderer));
+    resultToRender.put(type, result);
     return this;
   }
 
@@ -214,7 +236,7 @@ public ResultType call() throws Exception {
 
   @VisibleForTesting
   Renderer getRenderer(ResultType resultType, HttpServletRequest request) {
-    Set<Result> results = checkNotNull(resultToRender.get(resultType));
+    Collection<Result> results = checkNotNull(resultToRender.get(resultType));
     Result tempToUse = null;
     for (Result res : results) {
       if (tempToUse == null || res.isExactMimeTypeMatch(request.getHeader("Accept"))) {

java/server/src/org/openqa/selenium/remote/server/rest/UrlMapper.java

@@ -19,15 +19,18 @@
 
 import org.openqa.selenium.remote.server.DriverSessions;
 
-import java.util.LinkedHashMap;
+import com.google.common.collect.LinkedHashMultimap;
+import com.google.common.collect.Multimap;
+
+import java.util.Collection;
 import java.util.LinkedHashSet;
 import java.util.Map;
 import java.util.Set;
 import java.util.logging.Logger;
 
 public class UrlMapper {
 
-  private final Map<ResultType, Renderer> globals = new LinkedHashMap<ResultType, Renderer>();
+  private final Multimap<ResultType, Result> globals = LinkedHashMultimap.create();
   private final Set<ResultConfig> configs = new LinkedHashSet<ResultConfig>();
   private final DriverSessions sessions;
   private final Logger log;
@@ -45,8 +48,11 @@ public ResultConfig bind(String url, Class<? extends Handler> handlerClazz) {
 
     ResultConfig config = new ResultConfig(url, handlerClazz, sessions, log);
     configs.add(config);
-    for (Map.Entry<ResultType, Renderer> entry : globals.entrySet()) {
-      config.on(entry.getKey(), entry.getValue());
+    Map<ResultType, Collection<Result>> map = globals.asMap();
+    for (Map.Entry<ResultType, Collection<Result>> entry : map.entrySet()) {
+      for (Result result : entry.getValue()) {
+        config.on(entry.getKey(), result);
+      }
     }
     return config;
   }
@@ -61,11 +67,11 @@ public ResultConfig getConfig(String url) {
     return null;
   }
 
-  public void addGlobalHandler(ResultType type, Renderer renderer) {
-    globals.put(type, renderer);
+  public void addGlobalHandler(ResultType type, Result result) {
+    globals.put(type, result);
 
     for (ResultConfig config : configs) {
-      config.on(type, renderer);
+      config.on(type, result);
     }
   }
 }

java/server/test/org/openqa/selenium/remote/server/rest/UrlMapperTest.java

@@ -70,9 +70,10 @@ public void testShouldInjectDependenciesViaTheConstructor() throws Exception {
 
   public void testAppliesGlobalHandlersToNewConfigs() {
     Renderer renderer = new StubRenderer();
+    Result result = new Result("", renderer);
     HttpServletRequest mockRequest = context.mock(HttpServletRequest.class);
 
-    mapper.addGlobalHandler(ResultType.SUCCESS, renderer);
+    mapper.addGlobalHandler(ResultType.SUCCESS, result);
     mapper.bind("/example", SessionHandler.class);
 
     ResultConfig config = mapper.getConfig("/example");
@@ -81,10 +82,30 @@ public void testAppliesGlobalHandlersToNewConfigs() {
 
   public void testAppliesNewGlobalHandlersToExistingConfigs() {
     Renderer renderer = new StubRenderer();
+    Result result = new Result("", renderer);
     HttpServletRequest mockRequest = context.mock(HttpServletRequest.class);
 
     mapper.bind("/example", SessionHandler.class);
-    mapper.addGlobalHandler(ResultType.SUCCESS, renderer);
+    mapper.addGlobalHandler(ResultType.SUCCESS, result);
+
+    ResultConfig config = mapper.getConfig("/example");
+    assertEquals(renderer, config.getRenderer(ResultType.SUCCESS, mockRequest));
+  }
+
+  public void testPermitsMultipleGlobalHandlersWithDifferentMimeTypes() {
+    Renderer renderer = new StubRenderer();
+
+    final HttpServletRequest mockRequest = context.mock(HttpServletRequest.class);
+
+    context.checking(new Expectations() {{
+      allowing(mockRequest).getHeader("Accept");
+      will(returnValue("application/json"));
+    }});
+
+    mapper.addGlobalHandler(ResultType.SUCCESS, new Result("", new StubRenderer()));
+    mapper.addGlobalHandler(ResultType.SUCCESS, new Result("application/json", renderer ));
+    mapper.bind("/example", SessionHandler.class)
+        .on(ResultType.SUCCESS, new Result("text/plain", new StubRenderer()));
 
     ResultConfig config = mapper.getConfig("/example");
     assertEquals(renderer, config.getRenderer(ResultType.SUCCESS, mockRequest));