Fix tests failing because of ProtocolHandshake

Author: shs96c

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

@@ -217,10 +217,13 @@ public void writeTo(Appendable appendable) throws IOException {
     try (JsonOutput json = new Json().newOutput(appendable)) {
       json.beginObject();
 
-      @SuppressWarnings("unchecked")
-      Map<String, Object> first = (Map<String, Object>) stream().findFirst()
-          .orElse(new ImmutableCapabilities())
-          .asMap();
+      Map<String, Object> first = getOss();
+      if (first == null) {
+        //noinspection unchecked
+        first = (Map<String, Object>) stream().findFirst()
+            .orElse(new ImmutableCapabilities())
+            .asMap();
+      }
 
       // Write the first capability we get as the desired capability.
       json.name("desiredCapabilities");
@@ -295,7 +298,7 @@ public Stream<ImmutableCapabilities> stream() throws IOException {
   }
 
   public ImmutableSet<Dialect> getDownstreamDialects() {
-    return dialects.isEmpty() ? ImmutableSet.of(Dialect.OSS) : dialects;
+    return dialects.isEmpty() ? ImmutableSet.of(DEFAULT_DIALECT) : dialects;
   }
 
   @Override
@@ -325,7 +328,7 @@ private Stream<Map<String, Object>> getW3C() throws IOException {
     // then add magic to generate each of the w3c capabilities. For the sake of simplicity, we're
     // going to make the (probably wrong) assumption we can hold all of the firstMatch values and
     // alwaysMatch value in memory at the same time.
-    Map<String, Object> oss = getOss();
+    Map<String, Object> oss = convertOssToW3C(getOss());
     Stream<Map<String, Object>> fromOss;
     if (oss != null) {
       Set<String> usedKeys = new HashSet<>();
@@ -357,11 +360,10 @@ private Stream<Map<String, Object>> getW3C() throws IOException {
       // into the stream to form a unified set of capabilities. Woohoo!
       fromOss = firsts.stream()
           .map(first -> ImmutableMap.<String, Object>builder().putAll(always).putAll(first).build())
-          .map(this::convertOssToW3C)
+          .map(this::applyTransforms)
           .map(map -> map.entrySet().stream()
               .filter(entry -> ACCEPTED_W3C_PATTERNS.test(entry.getKey()))
              .collect(ImmutableMap.toImmutableMap(Map.Entry::getKey, Map.Entry::getValue)))
-          .filter(map -> !map.isEmpty())
           .map(map -> (Map<String, Object>) map);
     } else {
       fromOss = Stream.of();
@@ -386,10 +388,14 @@ private Stream<Map<String, Object>> getW3C() throws IOException {
           .map(first -> ImmutableMap.<String, Object>builder().putAll(always).putAll(first).build());
     }
 
-    return Stream.concat(fromOss, fromW3c);
+    return Stream.concat(fromOss, fromW3c).distinct();
   }
 
-  private Map<String, Object> convertOssToW3C(ImmutableMap<String, Object> capabilities) {
+  private Map<String, Object> convertOssToW3C(Map<String, Object> capabilities) {
+    if (capabilities == null) {
+      return null;
+    }
+
     Map<String, Object> toReturn = new TreeMap<>();
     toReturn.putAll(capabilities);
 

java/client/src/org/openqa/selenium/remote/session/ChromeFilter.java

@@ -18,22 +18,36 @@
 package org.openqa.selenium.remote.session;
 
 import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.ImmutableSortedMap;
+import com.google.common.collect.Ordering;
 
 import java.util.Map;
 import java.util.Objects;
+import java.util.TreeMap;
+import java.util.function.Supplier;
+import java.util.stream.Collectors;
 
 public class ChromeFilter implements CapabilitiesFilter {
   @Override
   public Map<String, Object> apply(Map<String, Object> unmodifiedCaps) {
-    ImmutableMap<String, Object> caps = unmodifiedCaps.entrySet().parallelStream()
+    Map<String, Object> caps = unmodifiedCaps.entrySet().parallelStream()
         .filter(
             entry ->
                 ("browserName".equals(entry.getKey()) && "chrome".equals(entry.getValue())) ||
                 entry.getKey().startsWith("goog:") ||
                 "chromeOptions".equals(entry.getKey()))
         .filter(entry -> Objects.nonNull(entry.getValue()))
         .distinct()
-        .collect(ImmutableMap.toImmutableMap(Map.Entry::getKey, Map.Entry::getValue));
+        .collect(Collectors.toMap(
+            Map.Entry::getKey,
+            Map.Entry::getValue,
+            (l, r) -> r,
+            TreeMap::new));
+
+    // We may need to map the chromeoptions to the new form
+    if (caps.containsKey("chromeOptions") && !caps.containsKey("goog:chromeOptions")) {
+      caps.put("goog:chromeOptions", caps.get("chromeOptions"));
+    }
 
     return caps.isEmpty() ? null : caps;
   }

java/client/src/org/openqa/selenium/remote/session/FirefoxFilter.java

@@ -17,24 +17,33 @@
 
 package org.openqa.selenium.remote.session;
 
-import com.google.common.collect.ImmutableMap;
-
 import java.util.Map;
 import java.util.Objects;
+import java.util.TreeMap;
+import java.util.stream.Collectors;
 
 public class FirefoxFilter implements CapabilitiesFilter {
   // Note: we don't take a dependency on the FirefoxDriver jar as it might not be on the classpath
 
   @Override
   public Map<String, Object> apply(Map<String, Object> unmodifiedCaps) {
-    ImmutableMap<String, Object> caps = unmodifiedCaps.entrySet().parallelStream()
+    Map<String, Object> caps = unmodifiedCaps.entrySet().parallelStream()
         .filter(entry ->
                     ("browserName".equals(entry.getKey()) && "firefox".equals(entry.getValue())) ||
                     entry.getKey().startsWith("firefox_") ||
                     entry.getKey().startsWith("moz:"))
         .filter(entry -> Objects.nonNull(entry.getValue()))
-        .distinct()
-        .collect(ImmutableMap.toImmutableMap(Map.Entry::getKey, Map.Entry::getValue));
+        .collect(Collectors.toMap(
+            Map.Entry::getKey,
+            Map.Entry::getValue,
+            (l, r) -> l,
+            TreeMap::new));
+
+    // If we only have marionette in the caps, the user is asking for firefox. Make sure we inject
+    // the browser name to be sure.
+    if (unmodifiedCaps.containsKey("marionette") && !caps.containsKey("browserName")) {
+      caps.put("browserName", "firefox");
+    }
 
     return caps.isEmpty() ? null : caps;
   }

java/client/test/org/openqa/selenium/remote/ProtocolHandshakeTest.java

@@ -217,8 +217,6 @@ public void shouldNotIncludeNonProtocolExtensionKeys() throws IOException {
 
     HttpRequest request = client.getRequest();
 
-    System.out.println(request.getContentString());
-
     Map<String, Object> handshakeRequest = new Gson().fromJson(
         request.getContentString(),
         new TypeToken<Map<String, Object>>() {}.getType());
@@ -267,20 +265,15 @@ public void firstMatchSeparatesCapsForDifferentBrowsers() throws IOException {
 
     List<Map<String, Object>> capabilities = mergeW3C(handshakeRequest);
 
-    System.out.println("capabilities = " + capabilities);
-
     assertThat(capabilities, containsInAnyOrder(
-        // OSS with valid w3c keys. It's fine that this won't match
-        ImmutableMap.of("browserName", "chrome", "moz:firefoxOptions", ImmutableMap.of()),
-        // Generated capabilities
         ImmutableMap.of("moz:firefoxOptions", ImmutableMap.of()),
         ImmutableMap.of("browserName", "chrome")));
   }
 
   @Test
   public void doesNotCreateFirstMatchForNonW3CCaps() throws IOException {
     Capabilities caps = new ImmutableCapabilities(
-        "chromeOptions", ImmutableMap.of(),
+        "cheese", ImmutableMap.of(),
         "moz:firefoxOptions", ImmutableMap.of(),
         "browserName", "firefox");
 
@@ -337,8 +330,6 @@ public void shouldLowerCaseProxyTypeForW3CRequest() throws IOException {
           assertEquals("autodetect", seenProxy.get("proxyType"));
         });
 
-    Object rawCaps = handshakeRequest.get("capabilities");
-
     Map<String, ?> jsonCaps = (Map<String, ?>) handshakeRequest.get("desiredCapabilities");
     Map<String, ?> seenProxy = (Map<String, ?>) jsonCaps.get("proxy");
     assertEquals("AUTODETECT", seenProxy.get("proxyType"));

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

@@ -157,12 +157,15 @@ public void ossPayloadWillBeFirstW3CPayload() {
 
     List<Capabilities> allCaps = getCapabilities(rawCapabilities);
 
-    assertEquals(2, allCaps.size());
+    assertEquals(3, allCaps.size());
     assertEquals(false, allCaps.get(0).getCapability("marionette"));
   }
 
   private List<Capabilities> getCapabilities(Map<String, Object> payload) {
     try (NewSessionPayload newSessionPayload = NewSessionPayload.create(payload)) {
+      StringBuilder b = new StringBuilder();
+      newSessionPayload.writeTo(b);
+      System.out.println("b = " + b);
       return newSessionPayload.stream().collect(ImmutableList.toImmutableList());
     } catch (IOException e) {
       throw new UncheckedIOException(e);