py: Modifies install_addon method to support installing unpacked addons (#10308)

gpt14 · web-flow · commit 8a23bcaa2d30 · 2022-03-01T14:57:43.000Z
* py: Modifies install_addon method to support installing unpacked addons

* Fix lint

* adds copy_file block to increase code readability

* Fix lint

* Fix lint
diff --git a/py/BUILD.bazel b/py/BUILD.bazel
@@ -65,6 +65,12 @@ copy_file(
     out = "selenium/webdriver/firefox/webdriver_prefs.json",
 )
 
+copy_file(
+    name = "favourite_color",
+    src = "//third_party/firebug:favourite_colour-1.1-an+fx.xpi",
+    out = "test/selenium/webdriver/firefox/favourite_colour-1.1-an+fx.xpi",
+)
+
 py_library(
     name = "selenium",
     srcs = glob(
@@ -78,6 +84,7 @@ py_library(
         ":get-attribute",
         ":is-displayed",
         ":mutation-listener",
+        ":favourite_color"
     ] + [":create-cdp-srcs-" + n for n in BROWSER_VERSIONS],
     imports = ["."],
     visibility = ["//visibility:public"],
diff --git a/py/selenium/webdriver/firefox/webdriver.py b/py/selenium/webdriver/firefox/webdriver.py
@@ -15,10 +15,13 @@
 # specific language governing permissions and limitations
 # under the License.
 import base64
+from io import BytesIO
+import os
 from shutil import rmtree
 import warnings
 from contextlib import contextmanager
 from typing import NoReturn
+import zipfile
 
 from selenium.webdriver.common.desired_capabilities import DesiredCapabilities
 from selenium.webdriver.remote.webdriver import WebDriver as RemoteWebDriver
@@ -232,7 +235,7 @@ def context(self, context):
         finally:
             self.set_context(initial_context)
 
-    def install_addon(self, path, temporary=None) -> str:
+    def install_addon(self, path, temporary=False) -> str:
         """
         Installs Firefox addon.
 
@@ -246,12 +249,21 @@ def install_addon(self, path, temporary=None) -> str:
 
                 driver.install_addon('/path/to/firebug.xpi')
         """
-        with open(path, 'rb') as file:
-            addon = (base64.b64encode(file.read()).decode('UTF-8'))
 
-        payload = {"addon": addon}
-        if temporary:
-            payload["temporary"] = temporary
+        if(os.path.isdir(path)):
+            fp = BytesIO()
+            path_root = len(path) + 1  # account for trailing slash
+            with zipfile.ZipFile(fp, 'w', zipfile.ZIP_DEFLATED) as zipped:
+                for base, dirs, files in os.walk(path):
+                    for fyle in files:
+                        filename = os.path.join(base, fyle)
+                        zipped.write(filename, filename[path_root:])
+            addon = base64.b64encode(fp.getvalue()).decode('UTF-8')
+        else:
+            with open(path, 'rb') as file:
+                addon = (base64.b64encode(file.read()).decode('UTF-8'))
+
+        payload = {"addon": addon, "temporary": temporary}
         return self.execute("INSTALL_ADDON", payload)["value"]
 
     def uninstall_addon(self, identifier) -> NoReturn:
diff --git a/py/test/selenium/webdriver/firefox/ff_installs_addons_tests.py b/py/test/selenium/webdriver/firefox/ff_installs_addons_tests.py
@@ -20,17 +20,25 @@
 from selenium.common.exceptions import WebDriverException
 
 
+def test_install_addon_temporary(driver):
+    extension = os.path.join(os.path.dirname(os.path.abspath(__file__)),
+                             'favourite_colour-1.1-an+fx.xpi')
+
+    id = driver.install_addon(extension, True)
+    assert id == 'favourite-colour-examples@mozilla.org'
+
+
 def test_install_addon(driver):
-    extension = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))),
-                             '../../../../third_party/firebug/favourite_colour-1.1-an+fx.xpi')
+    extension = os.path.join(os.path.dirname(os.path.abspath(__file__)),
+                             'favourite_colour-1.1-an+fx.xpi')
 
-    id = driver.install_addon(extension)
+    id = driver.install_addon(extension, False)
     assert id == 'favourite-colour-examples@mozilla.org'
 
 
 def test_uninstall_addon(driver):
-    extension = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))),
-                             '../../../../third_party/firebug/favourite_colour-1.1-an+fx.xpi')
+    extension = os.path.join(os.path.dirname(os.path.abspath(__file__)),
+                             'favourite_colour-1.1-an+fx.xpi')
 
     id = driver.install_addon(extension)
     try:
