Merge pull request #1452 from TaylorMcGinnis/handle_custom_desired_capabilities

Handle extension capabilities

Author: mdmintz

seleniumbase/core/browser_launcher.py

@@ -1068,6 +1068,7 @@ def get_remote_driver(
         screen_resolution = None
         browser_version = None
         platform_name = None
+        extension_capabilities = {}
         for key in desired_caps.keys():
             capabilities[key] = desired_caps[key]
             if key == "selenoid:options":
@@ -1079,6 +1080,8 @@ def get_remote_driver(
                 browser_version = desired_caps[key]
             elif key == "platform" or key == "platformName":
                 platform_name = desired_caps[key]
+            elif re.match("[a-zA-Z0-9]*:[a-zA-Z0-9]*", key):
+                extension_capabilities[key] = desired_caps[key]
         if selenium4:
             chrome_options.set_capability("cloud:options", capabilities)
             if selenoid:
@@ -1093,6 +1096,10 @@ def get_remote_driver(
             if platform_name:
                 plat_name = platform_name
                 chrome_options.set_capability("platformName", plat_name)
+            if extension_capabilities:
+                for key in extension_capabilities:
+                    ext_caps = extension_capabilities
+                    chrome_options.set_capability(key, ext_caps[key])
             return webdriver.Remote(
                 command_executor=address,
                 options=chrome_options,
@@ -1133,6 +1140,7 @@ def get_remote_driver(
         screen_resolution = None
         browser_version = None
         platform_name = None
+        extension_capabilities = {}
         for key in desired_caps.keys():
             capabilities[key] = desired_caps[key]
             if key == "selenoid:options":
@@ -1144,6 +1152,8 @@ def get_remote_driver(
                 browser_version = desired_caps[key]
             elif key == "platform" or key == "platformName":
                 platform_name = desired_caps[key]
+            elif re.match("[a-zA-Z0-9]*:[a-zA-Z0-9]*", key):
+                extension_capabilities[key] = desired_caps[key]
         if selenium4:
             firefox_options.set_capability("cloud:options", capabilities)
             if selenoid:
@@ -1158,6 +1168,10 @@ def get_remote_driver(
             if platform_name:
                 plat_name = platform_name
                 firefox_options.set_capability("platformName", plat_name)
+            if extension_capabilities:
+                for key in extension_capabilities:
+                    ext_caps = extension_capabilities
+                    firefox_options.set_capability(key, ext_caps[key])
             return webdriver.Remote(
                 command_executor=address,
                 options=firefox_options,
@@ -1278,6 +1292,7 @@ def get_remote_driver(
         screen_resolution = None
         browser_version = None
         platform_name = None
+        extension_capabilities = {}
         for key in desired_caps.keys():
             capabilities[key] = desired_caps[key]
             if key == "selenoid:options":
@@ -1289,6 +1304,8 @@ def get_remote_driver(
                 browser_version = desired_caps[key]
             elif key == "platform" or key == "platformName":
                 platform_name = desired_caps[key]
+            elif re.match("[a-zA-Z0-9]*:[a-zA-Z0-9]*", key):
+                extension_capabilities[key] = desired_caps[key]
         if selenium4:
             opera_options.set_capability("cloud:options", capabilities)
             if selenoid:
@@ -1303,6 +1320,10 @@ def get_remote_driver(
             if platform_name:
                 plat_name = platform_name
                 opera_options.set_capability("platformName", plat_name)
+            if extension_capabilities:
+                for key in extension_capabilities:
+                    ext_caps = extension_capabilities
+                    opera_options.set_capability(key, ext_caps[key])
             return webdriver.Remote(
                 command_executor=address,
                 options=opera_options,