Allow SessionFactory to indicate whether they support particular Capabilities

Also since the default new session pipeline includes all the
available and installed there's no need for a fallback to be
installed.

Author: shs96c

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

@@ -29,15 +29,15 @@
 import static org.openqa.selenium.remote.CapabilityType.BROWSER_NAME;
 
 import com.google.common.annotations.VisibleForTesting;
+import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 
 import org.openqa.selenium.Capabilities;
 import org.openqa.selenium.ImmutableCapabilities;
 import org.openqa.selenium.WebDriver;
 import org.openqa.selenium.remote.Dialect;
 
-import java.util.LinkedHashMap;
-import java.util.Map;
+import java.util.List;
 import java.util.Objects;
 import java.util.Optional;
 import java.util.ServiceLoader;
@@ -63,15 +63,15 @@ public class ActiveSessionFactory implements SessionFactory {
     }
   };
 
-  private volatile Map<Predicate<Capabilities>, SessionFactory> factories;
+  private volatile List<SessionFactory> factories;
 
   public ActiveSessionFactory() {
     // Insertion order matters. The first matching predicate is always used for matching.
-    Map<Predicate<Capabilities>, SessionFactory> builder = new LinkedHashMap<>();
+    ImmutableList.Builder<SessionFactory> builder = ImmutableList.builder();
 
     // Allow user-defined factories to override default ones.
     StreamSupport.stream(loadDriverProviders().spliterator(), false)
-        .forEach(p -> builder.put(p::canCreateDriverInstanceFor, new InMemorySession.Factory(p)));
+        .forEach(p -> builder.add(new InMemorySession.Factory(p)));
 
     ImmutableMap.<Predicate<Capabilities>, String>builder()
         .put(caps -> {
@@ -96,13 +96,13 @@ public ActiveSessionFactory() {
         .build()
         .entrySet().stream()
         .filter(e -> CLASS_EXISTS.apply(e.getValue()) != null)
-        .forEach(e -> builder.put(e.getKey(), new ServicedSession.Factory(e.getValue())));
+        .forEach(e -> builder.add(new ServicedSession.Factory(e.getKey(), e.getValue())));
 
     // Attempt to bind the htmlunitdriver if it's present.
     bind(builder, "org.openqa.selenium.htmlunit.HtmlUnitDriver", browserName(HTMLUNIT),
          new ImmutableCapabilities(BROWSER_NAME, HTMLUNIT));
 
-    this.factories = ImmutableMap.copyOf(builder);
+    this.factories = builder.build();
   }
 
   public synchronized ActiveSessionFactory bind(
@@ -113,11 +113,11 @@ public synchronized ActiveSessionFactory bind(
 
     LOG.info(String.format("Binding %s to respond to %s", useThis, onThis));
 
-    LinkedHashMap<Predicate<Capabilities>, SessionFactory> newMap = new LinkedHashMap<>();
-    newMap.put(onThis, useThis);
-    newMap.putAll(factories);
+    ImmutableList.Builder<SessionFactory> builder = ImmutableList.builder();
+    builder.add(useThis);
+    builder.addAll(factories);
 
-    factories = newMap;
+    factories = builder.build();
 
     return this;
   }
@@ -128,7 +128,7 @@ protected Iterable<DriverProvider> loadDriverProviders() {
   }
 
   private void bind(
-      Map<Predicate<Capabilities>, SessionFactory> builder,
+      ImmutableList.Builder<SessionFactory> builder,
       String className,
       Predicate<Capabilities> predicate,
       Capabilities capabilities) {
@@ -139,9 +139,7 @@ private void bind(
       }
 
       Class<? extends WebDriver> driverClass = clazz.asSubclass(WebDriver.class);
-      builder.put(
-          predicate,
-          new InMemorySession.Factory(new DefaultDriverProvider(capabilities, driverClass)));
+      builder.add(new InMemorySession.Factory(new DefaultDriverProvider(capabilities, driverClass)));
     } catch (ClassCastException ignored) {
       // Just carry on. Everything is fine.
     }
@@ -161,13 +159,20 @@ private static Predicate<Capabilities> containsKey(Pattern pattern) {
     return toCompare -> toCompare.asMap().keySet().stream().anyMatch(pattern.asPredicate());
   }
 
+  @Override
+  public boolean isSupporting(Capabilities capabilities) {
+    return factories.stream()
+        .map(factory -> factory.isSupporting(capabilities))
+        .reduce(Boolean::logicalOr)
+        .orElse(false);
+  }
+
   @Override
   public Optional<ActiveSession> apply(Set<Dialect> downstreamDialects, Capabilities caps) {
     LOG.info("Capabilities are: " + caps);
-    return factories.entrySet().stream()
-        .filter(e -> e.getKey().test(caps))
-        .peek(e -> LOG.info(String.format("%s matched %s", caps, e.getValue())))
-        .map(Map.Entry::getValue)
+    return factories.stream()
+        .filter(factory -> factory.isSupporting(caps))
+        .peek(factory -> LOG.info(String.format("%s matched %s", caps, factory)))
         .map(factory -> factory.apply(downstreamDialects, caps))
         .filter(Optional::isPresent)
         .map(Optional::get)

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

@@ -17,11 +17,7 @@
 
 package org.openqa.selenium.remote.server;
 
-import org.openqa.selenium.remote.service.DriverService;
-
-import java.util.Optional;
 import java.util.logging.Logger;
-import java.util.stream.Stream;
 
 /**
  * Used to represent the {@link NewSessionPipeline} that is typically used in the
@@ -35,34 +31,8 @@ private DefaultPipeline() {
     // Utility class
   }
 
-  public static NewSessionPipeline.Builder createPipelineWithDefaultFallbacks() {
-    // Set up the pipeline to inject
-    SessionFactory fallback = Stream.of(
-        "org.openqa.selenium.chrome.ChromeDriverService",
-        "org.openqa.selenium.firefox.GeckoDriverService",
-        "org.openqa.selenium.edge.EdgeDriverService",
-        "org.openqa.selenium.ie.InternetExplorerDriverService",
-        "org.openqa.selenium.safari.SafariDriverService")
-        .filter(name -> {
-          try {
-            Class.forName(name).asSubclass(DriverService.class);
-            return true;
-          } catch (ReflectiveOperationException e) {
-            return false;
-          }
-        })
-        .findFirst()
-        .map(serviceName -> {
-          SessionFactory factory = new ServicedSession.Factory(serviceName);
-          return (SessionFactory) (dialects, caps) -> {
-            LOG.info("Using default factory: " + serviceName);
-            return factory.apply(dialects, caps);
-          };
-        })
-        .orElse((dialects, caps) -> Optional.empty());
-
+  public static NewSessionPipeline.Builder createDefaultPipeline() {
     return NewSessionPipeline.builder()
-        .add(new ActiveSessionFactory())
-        .fallback(fallback);
+        .add(new ActiveSessionFactory());
   }
 }

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

@@ -125,6 +125,12 @@ public Factory(DriverProvider provider) {
       this.provider = provider;
     }
 
+
+    @Override
+    public boolean isSupporting(Capabilities capabilities) {
+      return provider.canCreateDriverInstanceFor(capabilities);
+    }
+
     @Override
     public Optional<ActiveSession> apply(Set<Dialect> downstreamDialects, Capabilities caps) {
       // Assume the blob fits in the available memory.

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

@@ -2,15 +2,18 @@
 
 import com.google.common.collect.ImmutableList;
 
+import org.openqa.selenium.Capabilities;
 import org.openqa.selenium.ImmutableCapabilities;
 import org.openqa.selenium.SessionNotCreatedException;
+import org.openqa.selenium.remote.Dialect;
 import org.openqa.selenium.remote.NewSessionPayload;
 
 import java.io.IOException;
 import java.util.LinkedList;
 import java.util.List;
 import java.util.Objects;
 import java.util.Optional;
+import java.util.Set;
 import java.util.function.Function;
 
 public class NewSessionPipeline {
@@ -41,6 +44,7 @@ public ActiveSession createNewSession(NewSessionPayload payload) throws IOExcept
           return caps;
         })
         .map(caps -> factories.stream()
+            .filter(factory -> factory.isSupporting(caps))
             .map(factory -> factory.apply(payload.getDownstreamDialects(), caps))
             .filter(Optional::isPresent)
             .map(Optional::get)
@@ -57,7 +61,17 @@ public ActiveSession createNewSession(NewSessionPayload payload) throws IOExcept
 
   public static class Builder {
     private List<SessionFactory> factories = new LinkedList<>();
-    private SessionFactory fallback = (dialects, caps) -> Optional.empty();
+    private SessionFactory fallback = new SessionFactory() {
+      @Override
+      public boolean isSupporting(Capabilities capabilities) {
+        return false;
+      }
+
+      @Override
+      public Optional<ActiveSession> apply(Set<Dialect> downstreamDialects, Capabilities capabilities) {
+        return Optional.empty();
+      }
+    };
     private List<Function<ImmutableCapabilities, ImmutableCapabilities>> mutators = new LinkedList<>();
 
     private Builder() {

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

@@ -201,7 +201,7 @@ public boolean boot() {
   }
 
   private NewSessionPipeline createPipeline(StandaloneConfiguration configuration) {
-    NewSessionPipeline.Builder builder = DefaultPipeline.createPipelineWithDefaultFallbacks();
+    NewSessionPipeline.Builder builder = DefaultPipeline.createDefaultPipeline();
 
     if (configuration instanceof GridNodeConfiguration) {
       ((GridNodeConfiguration) configuration).capabilities.forEach(

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

@@ -39,6 +39,7 @@
 import java.util.Optional;
 import java.util.Set;
 import java.util.function.Function;
+import java.util.function.Predicate;
 import java.util.logging.Level;
 
 import javax.management.MalformedObjectNameException;
@@ -79,10 +80,13 @@ public void stop() {
 
   public static class Factory extends RemoteSession.Factory<DriverService> {
 
+    private final Predicate<Capabilities> key;
     private final Function<Capabilities, ? extends DriverService> createService;
     private final String serviceClassName;
 
-    public Factory(String serviceClassName) {
+    public Factory(Predicate<Capabilities> key, String serviceClassName) {
+      this.key = key;
+
       this.serviceClassName = serviceClassName;
       try {
         Class<? extends DriverService> driverClazz =
@@ -129,6 +133,11 @@ public Factory(String serviceClassName) {
       }
     }
 
+    @Override
+    public boolean isSupporting(Capabilities capabilities) {
+      return key.test(capabilities);
+    }
+
     @Override
     public Optional<ActiveSession> apply(
         Set<Dialect> downstreamDialects,

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

@@ -25,5 +25,7 @@
 import java.util.Set;
 
 public interface SessionFactory {
+  boolean isSupporting(Capabilities capabilities);
+
   Optional<ActiveSession> apply(Set<Dialect> downstreamDialects, Capabilities capabilities);
 }

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

@@ -86,7 +86,7 @@ public void init() {
     NewSessionPipeline pipeline =
         (NewSessionPipeline) getServletContext().getAttribute(NEW_SESSION_PIPELINE_KEY);
     if (pipeline == null) {
-      pipeline = DefaultPipeline.createPipelineWithDefaultFallbacks().create();
+      pipeline = DefaultPipeline.createDefaultPipeline().create();
       getServletContext().setAttribute(NEW_SESSION_PIPELINE_KEY, pipeline);
     }
 

java/server/test/org/openqa/selenium/remote/server/ActiveSessionFactoryTest.java

@@ -31,6 +31,7 @@
 
 import java.io.IOException;
 import java.util.Optional;
+import java.util.Set;
 
 public class ActiveSessionFactoryTest {
 
@@ -56,7 +57,20 @@ public void canBindNewFactoriesAtRunTime() throws IOException {
     ActiveSession session = Mockito.mock(ActiveSession.class);
 
     ActiveSessionFactory sessionFactory = new ActiveSessionFactory()
-        .bind(caps -> "cheese".equals(caps.getBrowserName()), (dialects, caps) -> Optional.of(session));
+        .bind(caps ->
+                  "cheese".equals(caps.getBrowserName()),
+              new SessionFactory() {
+                @Override
+                public boolean isSupporting(Capabilities capabilities) {
+                  return true;
+                }
+
+                @Override
+                public Optional<ActiveSession> apply(Set<Dialect> downstreamDialects,
+                                                     Capabilities capabilities) {
+                  return Optional.of(session);
+                }
+              });
 
     ActiveSession created = sessionFactory.apply(ImmutableSet.copyOf(Dialect.values()), toPayload("cheese")).get();
 

java/server/test/org/openqa/selenium/remote/server/NewSessionPipelineTest.java

@@ -17,8 +17,10 @@
 
 package org.openqa.selenium.remote.server;
 
+import static org.junit.Assert.assertEquals;
 import static org.mockito.Mockito.any;
 import static org.mockito.Mockito.argThat;
+import static org.mockito.Mockito.atLeast;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.verifyZeroInteractions;
@@ -39,7 +41,9 @@ public class NewSessionPipelineTest {
   @Test
   public void shouldCallSessionFactory() throws IOException {
     SessionFactory factory = mock(SessionFactory.class);
+    when(factory.isSupporting(any())).thenReturn(true);
     SessionFactory fallback = mock(SessionFactory.class);
+    when(fallback.isSupporting(any())).thenReturn(true);
     ActiveSession session = mock(ActiveSession.class);
     when(factory.apply(any(), any())).thenReturn(Optional.of(session));
 
@@ -77,6 +81,7 @@ public void shouldBeAbleToFallBack() throws IOException {
   @Test
   public void shouldUseMutators() throws IOException {
     SessionFactory factory = mock(SessionFactory.class);
+    when(factory.isSupporting(any())).thenReturn(true);
     ActiveSession session = mock(ActiveSession.class);
     when(factory.apply(any(), any())).thenReturn(Optional.of(session));
     FirefoxMutator mutator = new FirefoxMutator(new ImmutableCapabilities(
@@ -89,9 +94,34 @@ public void shouldUseMutators() throws IOException {
             "alwaysMatch", ImmutableMap.of("browserName", "firefox")));
 
     NewSessionPipeline pipeline = NewSessionPipeline.builder()
-        .add(factory).addCapabilitiesMutator(mutator).create();
+        .add(factory)
+        .addCapabilitiesMutator(mutator)
+        .create();
 
     pipeline.createNewSession(NewSessionPayload.create(caps));
     verify(factory).apply(any(), argThat(cap -> cap.getCapability("marionette").equals(true)));
   }
+
+  @Test
+  public void shouldNotUseFactoriesThatDoNotSupportTheCapabilities() throws IOException {
+    SessionFactory toBeIgnored = mock(SessionFactory.class);
+    when(toBeIgnored.isSupporting(any())).thenReturn(false);
+    when(toBeIgnored.apply(any(), any())).thenThrow(new AssertionError("Must not be called"));
+
+    ActiveSession session = mock(ActiveSession.class);
+    SessionFactory toBeUsed = mock(SessionFactory.class);
+    when(toBeUsed.isSupporting(any())).thenReturn(true);
+    when(toBeUsed.apply(any(), any())).thenReturn(Optional.of(session));
+
+    NewSessionPipeline pipeline = NewSessionPipeline.builder()
+        .add(toBeIgnored)
+        .add(toBeUsed)
+        .create();
+
+    ActiveSession seen =
+        pipeline.createNewSession(NewSessionPayload.create(new ImmutableCapabilities()));
+
+    assertEquals(session, seen);
+    verify(toBeIgnored, atLeast(1)).isSupporting(any());
+  }
 }