use FileDetector to install Firefox addons if one is set on the driver

SeleniumHQ · Sep 29, 2021 · 5708473 · 5708473
1 parent 60b2cff
commit 5708473
Show file tree

Hide file tree

Showing 3 changed files with 45 additions and 6 deletions.
diff --git a/java/src/org/openqa/selenium/firefox/AddHasExtensions.java b/java/src/org/openqa/selenium/firefox/AddHasExtensions.java
@@ -20,18 +20,26 @@
 import com.google.auto.service.AutoService;
 import com.google.common.collect.ImmutableMap;
 import org.openqa.selenium.Capabilities;
+import org.openqa.selenium.WebDriver;
+import org.openqa.selenium.WebDriverException;
+import org.openqa.selenium.WrapsDriver;
 import org.openqa.selenium.internal.Require;
+import org.openqa.selenium.io.Zip;
 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.RemoteWebDriver;
 import org.openqa.selenium.remote.http.HttpMethod;
 
+import java.io.File;
+import java.io.IOException;
 import java.nio.file.Path;
 import java.util.Map;
 import java.util.function.Predicate;
 
 import static org.openqa.selenium.remote.Browser.FIREFOX;
+import static org.openqa.selenium.remote.DriverCommand.UPLOAD_FILE;
 
 @AutoService({AdditionalHttpCommands.class, AugmenterProvider.class})
 public class AddHasExtensions implements AugmenterProvider<HasExtensions>, AdditionalHttpCommands {
@@ -65,9 +73,23 @@ public HasExtensions getImplementation(Capabilities capabilities, ExecuteMethod
       public String installExtension(Path path) {
         Require.nonNull("Extension Path", path);
 
-        return (String) executeMethod.execute(
-          INSTALL_EXTENSION,
-          ImmutableMap.of("path", path.toAbsolutePath().toString(), "temporary", false));
+        if (executeMethod instanceof WrapsDriver) {
+          WebDriver wrapped = ((WrapsDriver) executeMethod).getWrappedDriver();
+          if (wrapped instanceof RemoteWebDriver) {
+            File localFile = ((RemoteWebDriver) wrapped).getFileDetector().getLocalFile(path.toString());
+            if (localFile != null) {
+              try {
+                String zip = Zip.zip(localFile);
+                String newPath = (String) executeMethod.execute(UPLOAD_FILE, ImmutableMap.of("file", zip));
+                return installExtensionAtPath(newPath);
+              } catch (IOException e) {
+                throw new WebDriverException("Cannot upload " + localFile, e);
+              }
+            }
+          }
+        }
+
+        return installExtensionAtPath(path.toString());
       }
 
       @Override
@@ -76,6 +98,12 @@ public void uninstallExtension(String extensionId) {
 
         executeMethod.execute(UNINSTALL_EXTENSION, ImmutableMap.of("id", extensionId));
       }
+
+      private String installExtensionAtPath(String path) {
+        return (String) executeMethod.execute(
+          INSTALL_EXTENSION,
+          ImmutableMap.of("path", path, "temporary", false));
+      }
     };
   }
 }
diff --git a/java/src/org/openqa/selenium/remote/RemoteExecuteMethod.java b/java/src/org/openqa/selenium/remote/RemoteExecuteMethod.java
@@ -17,11 +17,13 @@
 
 package org.openqa.selenium.remote;
 
+import org.openqa.selenium.WebDriver;
+import org.openqa.selenium.WrapsDriver;
 import org.openqa.selenium.internal.Require;
 
 import java.util.Map;
 
-public class RemoteExecuteMethod implements ExecuteMethod {
+public class RemoteExecuteMethod implements ExecuteMethod, WrapsDriver {
   private final RemoteWebDriver driver;
 
   public RemoteExecuteMethod(RemoteWebDriver driver) {
@@ -40,4 +42,8 @@ public Object execute(String commandName, Map<String, ?> parameters) {
 
     return response.getValue();
   }
+
+  @Override public WebDriver getWrappedDriver() {
+    return this.driver;
+  }
 }
diff --git a/java/test/org/openqa/selenium/firefox/FirefoxDriverTest.java b/java/test/org/openqa/selenium/firefox/FirefoxDriverTest.java
@@ -37,6 +37,7 @@
 import org.openqa.selenium.remote.Command;
 import org.openqa.selenium.remote.CommandExecutor;
 import org.openqa.selenium.remote.DriverCommand;
+import org.openqa.selenium.remote.LocalFileDetector;
 import org.openqa.selenium.remote.RemoteWebDriver;
 import org.openqa.selenium.remote.SessionId;
 import org.openqa.selenium.remote.UnreachableBrowserException;
@@ -503,8 +504,12 @@ public void testFirefoxCanNativelyClickOverlappingElements() {
   public void canAddRemoveExtensions() {
     Path extension = InProject.locate("third_party/firebug/favourite_colour-1.1-an+fx.xpi");
 
-    ((HasExtensions) driver).installExtension(extension);
-    ((HasExtensions) driver).uninstallExtension("[email protected]");
+    if (driver.getClass().equals(RemoteWebDriver.class)) {
+      ((RemoteWebDriver) driver).setFileDetector(new LocalFileDetector());
+    }
+
+    String id = ((HasExtensions) driver).installExtension(extension);
+    ((HasExtensions) driver).uninstallExtension(id);
   }
 
   @Test