Implement Firefox endpoint for setting context and allow to be Augmented

java/src/org/openqa/selenium/firefox/AddHasContext.java

@@ -0,0 +1,67 @@
+// Licensed to the Software Freedom Conservancy (SFC) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The SFC licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+package org.openqa.selenium.firefox;
+
+import com.google.auto.service.AutoService;
+import com.google.common.collect.ImmutableMap;
+import org.openqa.selenium.Capabilities;
+import org.openqa.selenium.remote.AdditionalHttpCommands;
+import org.openqa.selenium.remote.AugmenterProvider;
+import org.openqa.selenium.remote.CommandInfo;
+import org.openqa.selenium.remote.ExecuteMethod;
+import org.openqa.selenium.remote.http.HttpMethod;
+
+import java.util.Map;
+import java.util.function.Predicate;
+
+import static org.openqa.selenium.remote.BrowserType.FIREFOX;
+
+@AutoService({AdditionalHttpCommands.class, AugmenterProvider.class})
+public class AddHasContext implements AugmenterProvider<HasContext>, AdditionalHttpCommands {
+
+  public static final String CONTEXT = "context";
+
+  private static final Map<String, CommandInfo> COMMANDS = ImmutableMap.of(
+    CONTEXT, new CommandInfo("/session/:sessionId/moz/context", HttpMethod.POST));
+
+  @Override
+  public Map<String, CommandInfo> getAdditionalCommands() {
+    return COMMANDS;
+  }
+
+  @Override
+  public Predicate<Capabilities> isApplicable() {
+    return caps -> FIREFOX.equals(caps.getBrowserName());
+  }
+
+  @Override
+  public Class<HasContext> getDescribedInterface() {
+    return HasContext.class;
+  }
+
+  @Override
+  public HasContext getImplementation(Capabilities capabilities, ExecuteMethod executeMethod) {
+    return new HasContext() {
+      @Override public void setContext(FirefoxCommandContext context) {
+        executeMethod.execute(
+          CONTEXT,
+          ImmutableMap.of(CONTEXT, context));
+      }
+    };
+  }
+}

java/src/org/openqa/selenium/firefox/FirefoxCommandContext.java

@@ -0,0 +1,37 @@
+// Licensed to the Software Freedom Conservancy (SFC) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The SFC licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+package org.openqa.selenium.firefox;
+
+/**
+ * Represents the valid values for the command context used for executing Firefox driver commands.
+ */
+public enum FirefoxCommandContext {
+  CONTENT("content"),
+  CHROME("chrome");
+
+  private String text;
+
+  FirefoxCommandContext(String text) {
+    this.text = text;
+  }
+
+  @Override
+  public String toString() {
+    return String.valueOf(text);
+  }
+}

java/src/org/openqa/selenium/firefox/FirefoxDriver.java

@@ -74,7 +74,7 @@
  * </pre>
  */
 public class FirefoxDriver extends RemoteWebDriver
-  implements WebStorage, HasExtensions, HasFullPageScreenshot, HasDevTools {
+  implements WebStorage, HasExtensions, HasFullPageScreenshot, HasContext, HasDevTools {
 
   public static final class SystemProperty {
 
@@ -155,6 +155,7 @@ private static Map<String, CommandInfo> getExtraCommands() {
   private final RemoteWebStorage webStorage;
   private final HasExtensions extensions;
   private final HasFullPageScreenshot fullPageScreenshot;
+  private final HasContext context;
   private final Optional<URI> cdpUri;
   private DevTools devTools;
 
@@ -197,6 +198,7 @@ private FirefoxDriver(FirefoxDriverCommandExecutor executor, FirefoxOptions opti
     webStorage = new RemoteWebStorage(getExecuteMethod());
     extensions = new AddHasExtensions().getImplementation(getCapabilities(), getExecuteMethod());
     fullPageScreenshot = new AddHasFullPageScreenshot().getImplementation(getCapabilities(), getExecuteMethod());
+    context = new AddHasContext().getImplementation(getCapabilities(), getExecuteMethod());
 
     Capabilities capabilities = super.getCapabilities();
     HttpClient.Factory clientFactory = HttpClient.Factory.createDefault();
@@ -285,6 +287,10 @@ public <X> X getFullPageScreenshotAs(OutputType<X> outputType) throws WebDriverE
     return fullPageScreenshot.getFullPageScreenshotAs(outputType);
   }
 
+  @Override public void setContext(FirefoxCommandContext commandContext) {
+    context.setContext(commandContext);
+  }
+
   private static Boolean forceMarionetteFromSystemProperty() {
     String useMarionette = System.getProperty(SystemProperty.DRIVER_USE_MARIONETTE);
     if (useMarionette == null) {

java/src/org/openqa/selenium/firefox/HasContext.java

@@ -0,0 +1,36 @@
+// Licensed to the Software Freedom Conservancy (SFC) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The SFC licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+package org.openqa.selenium.firefox;
+
+import org.openqa.selenium.Beta;
+
+/**
+ * Used by classes to indicate that they can install and uninstall browser extensions on the fly.
+ */
+@Beta
+public interface HasContext {
+
+  /**
+   * Context commands are operating on.
+   *
+   * @param context {@link FirefoxCommandContext} operating on page loaded in the browser or on browser elements hosting the page.
+   *
+   */
+  void setContext(FirefoxCommandContext context);
+
+}

java/test/org/openqa/selenium/firefox/FirefoxDriverTest.java

@@ -600,6 +600,60 @@ public void shouldAllowDriverToBeAugmentedWithFullHasPageScreenshot() throws Mal
     }
   }
 
+  @Test
+  public void shouldSetContext() {
+    FirefoxOptions options = new FirefoxOptions();
+    String dir = "foo/bar";
+    options.addPreference("browser.download.dir", dir);
+
+    FirefoxDriver driver = new FirefoxDriver(options);
+
+    try {
+      driver.setContext(FirefoxCommandContext.CHROME);
+      String result = (String) driver.executeScript("return Services.prefs.getStringPref('browser.download.dir')");
+      assertThat(result).isEqualTo(dir);
+    } finally {
+      driver.quit();
+    }
+  }
+
+  @Test
+  public void shouldAllowRemoteWebDriverBuilderToUseHasContext() throws MalformedURLException {
+    FirefoxOptions options = new FirefoxOptions();
+    String dir = "foo/bar";
+    options.addPreference("browser.download.dir", dir);
+    WebDriver driver = new RemoteWebDriver(new URL("http://localhost:4444/"), options);
+    WebDriver augmentedDriver = new Augmenter().augment(driver);
+
+    try {
+      ((HasContext) augmentedDriver).setContext(FirefoxCommandContext.CHROME);
+      String result = (String) ((JavascriptExecutor) driver).executeScript("return Services.prefs.getStringPref('browser.download.dir')");
+      assertThat(result).isEqualTo(dir);
+    } finally {
+      driver.quit();
+    }
+  }
+
+  @Test
+  public void shouldAllowRemoteWebDriverToAugmenterHasContext() {
+    FirefoxOptions options = new FirefoxOptions();
+    String dir = "foo/bar";
+    options.addPreference("browser.download.dir", dir);
+
+    WebDriver driver = RemoteWebDriver.builder()
+      .oneOf(options)
+      .address("http://localhost:4444/")
+      .build();
+
+    try {
+      ((HasContext) driver).setContext(FirefoxCommandContext.CHROME);
+      String result = (String) ((JavascriptExecutor) driver).executeScript("return Services.prefs.getStringPref('browser.download.dir')");
+      assertThat(result).isEqualTo(dir);
+    } finally {
+      driver.quit();
+    }
+  }
+
   private static class CustomFirefoxProfile extends FirefoxProfile {}
 
 }