Allow the new session pipeline capabilities to be amended

Author: shs96c

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

@@ -399,7 +399,7 @@ private static Map<String, Object> toMap(Object obj) {
         .collect(ImmutableMap.toImmutableMap(e -> String.valueOf(e.getKey()), Map.Entry::getValue));
   }
 
-  public Stream<Capabilities> stream() throws IOException {
+  public Stream<ImmutableCapabilities> stream() throws IOException {
     Stream<? extends Map<String, Object>> mapStream;
 
     if (getDownstreamDialects().contains(Dialect.W3C)) {

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

@@ -10,15 +10,21 @@
 import java.util.List;
 import java.util.Objects;
 import java.util.Optional;
+import java.util.function.Function;
 
 public class NewSessionPipeline {
 
   private final List<SessionFactory> factories;
   private final SessionFactory fallback;
+  private final List<Function<ImmutableCapabilities, ImmutableCapabilities>> mutators;
 
-  private NewSessionPipeline(List<SessionFactory> factories, SessionFactory fallback) {
+  private NewSessionPipeline(
+      List<SessionFactory> factories,
+      SessionFactory fallback,
+      List<Function<ImmutableCapabilities, ImmutableCapabilities>> mutators) {
     this.factories = factories;
     this.fallback = fallback;
+    this.mutators = mutators;
   }
 
   public static Builder builder() {
@@ -27,6 +33,12 @@ public static Builder builder() {
 
   public ActiveSession createNewSession(NewSessionPayload payload) throws IOException {
     return payload.stream()
+        .map(caps -> {
+          for (Function<ImmutableCapabilities, ImmutableCapabilities> mutator : mutators) {
+            caps = mutator.apply(caps);
+          }
+          return caps;
+        })
         .map(caps -> factories.stream()
             .map(factory -> factory.apply(payload.getDownstreamDialects(), caps))
             .filter(Optional::isPresent)
@@ -45,26 +57,30 @@ 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 List<Function<ImmutableCapabilities, ImmutableCapabilities>> mutators = new LinkedList<>();
 
     private Builder() {
       // Private class
     }
 
     public Builder add(SessionFactory factory) {
       factories.add(Objects.requireNonNull(factory, "Factory must not be null"));
-//      for (SessionFactory fac : more) {
-//        factories.add(Objects.requireNonNull(fac, "Factory must not be null"));
-//      }
       return this;
     }
 
     public Builder fallback(SessionFactory factory) {
-      fallback = factory;
+      fallback = Objects.requireNonNull(factory, "Fallback must not be null");
+      return this;
+    }
+
+    public Builder addCapabilitiesMutator(
+        Function<ImmutableCapabilities, ImmutableCapabilities> mutator) {
+      mutators.add(Objects.requireNonNull(mutator, "Mutator must not be null"));
       return this;
     }
 
     public NewSessionPipeline create() {
-      return new NewSessionPipeline(ImmutableList.copyOf(factories), fallback);
+      return new NewSessionPipeline(ImmutableList.copyOf(factories), fallback, mutators);
     }
   }
 }