[java] Support for sending firefox addon directory as temporary in remote sessions

TamsilAmani · web-flow · commit 89710bbf41c3 · 2022-04-22T12:43:24.000+05:30
diff --git a/java/src/org/openqa/selenium/firefox/AddHasExtensions.java b/java/src/org/openqa/selenium/firefox/AddHasExtensions.java
@@ -28,12 +28,19 @@
 import org.openqa.selenium.remote.ExecuteMethod;
 import org.openqa.selenium.remote.http.HttpMethod;
 
+import java.io.FileOutputStream;
 import java.io.IOException;
+import java.nio.file.FileVisitResult;
 import java.nio.file.Files;
 import java.nio.file.Path;
+import java.nio.file.Paths;
+import java.nio.file.SimpleFileVisitor;
+import java.nio.file.attribute.BasicFileAttributes;
 import java.util.Base64;
 import java.util.Map;
 import java.util.function.Predicate;
+import java.util.zip.ZipEntry;
+import java.util.zip.ZipOutputStream;
 
 import static org.openqa.selenium.remote.Browser.FIREFOX;
 
@@ -77,7 +84,12 @@ public String installExtension(Path path, Boolean temporary) {
 
         String encoded;
         try {
-          encoded = Base64.getEncoder().encodeToString(Files.readAllBytes(path));
+          if (Files.isDirectory(path)) {
+            encoded = Base64.getEncoder().encodeToString(Files.readAllBytes(zipDirectory(path)));
+          }
+          else {
+            encoded = Base64.getEncoder().encodeToString(Files.readAllBytes(path));
+          }
         } catch (IOException e) {
           throw new InvalidArgumentException(path + " is an invalid path", e);
         }
@@ -87,6 +99,22 @@ public String installExtension(Path path, Boolean temporary) {
           ImmutableMap.of("addon", encoded, "temporary", temporary));
       }
 
+      private Path zipDirectory(Path path) throws IOException {
+        Path extZip = Paths.get(path.getFileName().toString()+".zip");
+        try (ZipOutputStream zos = new ZipOutputStream(new FileOutputStream(extZip.toFile()))) {
+          Files.walkFileTree(path, new SimpleFileVisitor<Path>() {
+            @Override
+            public FileVisitResult visitFile(Path file, BasicFileAttributes attrs) throws IOException {
+              zos.putNextEntry(new ZipEntry(path.relativize(file).toString()));
+              Files.copy(file, zos);
+              zos.closeEntry();
+              return FileVisitResult.CONTINUE;
+            }
+          });
+        }
+        return extZip;
+      }
+
       @Override
       public void uninstallExtension(String extensionId) {
         Require.nonNull("Extension ID", extensionId);
diff --git a/java/test/org/openqa/selenium/firefox/FirefoxDriverTest.java b/java/test/org/openqa/selenium/firefox/FirefoxDriverTest.java
@@ -78,6 +78,7 @@
 public class FirefoxDriverTest extends JUnit4TestBase {
 
   private static final String EXT_PATH = "common/extensions/webextensions-selenium-example.xpi";
+  private static final String EXT_PATH_DIR = "common/extensions/webextensions-selenium-example";
   private static char[] CHARS =
       "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz1234567890!\"§$%&/()+*~#',.-_:;\\"
           .toCharArray();
@@ -556,6 +557,20 @@ public void canAddRemoveTempExtensions() {
     }
   }
 
+  @Test
+  public void canAddRemoveTempExtensionsDirectory() {
+    Path extension = InProject.locate(EXT_PATH_DIR);
+
+    String id = ((HasExtensions) driver).installExtension(extension, true);
+    assertThat(id).isEqualTo("webextensions-selenium-example@example.com");
+
+    try {
+      ((HasExtensions) driver).uninstallExtension(id);
+    } catch (WebDriverException ex) {
+      fail(ex.getMessage());
+    }
+  }
+
   @Test
   public void canTakeFullPageScreenshot() {
     File tempFile = ((HasFullPageScreenshot) driver).getFullPageScreenshotAs(OutputType.FILE);
