Add a command line option to set a proxy server for tests to use.

Author: mdmintz

seleniumbase/config/proxy_list.py

@@ -0,0 +1,28 @@
+"""
+Proxy Server "Phone Book".
+
+Simplify running browser tests through a proxy server
+by adding your frequently-used proxies here.
+
+Now you can do something like this on the command line:
+"pytest SOME_TEST.py --proxy=proxy1"
+
+Format of PROXY_LIST server entries:
+* "ip_address:port"
+
+Example proxies in PROXY_LIST below are not guaranteed to be active or secure.
+If you don't already have a proxy server to connect to,
+you can try finding one from one of following sites:
+* https://hidemy.name/en/proxy-list/?country=US&type=h#list
+* http://proxyservers.pro/proxy/list/protocol/http/country/US/
+"""
+
+PROXY_LIST = {
+    # "example1": "35.196.26.166:3128",  # (Example) - set your own proxy here
+    # "example2": "208.95.62.81:3128",  # (Example) - set your own proxy here
+    "proxy1": None,
+    "proxy2": None,
+    "proxy3": None,
+    "proxy4": None,
+    "proxy5": None,
+}

seleniumbase/core/browser_launcher.py

@@ -1,16 +1,26 @@
+import re
 import warnings
 from selenium import webdriver
 from selenium.common.exceptions import WebDriverException
 from selenium.webdriver.common.desired_capabilities import DesiredCapabilities
 from seleniumbase.config import settings
+from seleniumbase.config import proxy_list
 from seleniumbase.core import download_helper
 from seleniumbase.fixtures import constants
 
 
-def _create_firefox_profile(downloads_path):
+def _create_firefox_profile(downloads_path, proxy_string):
     profile = webdriver.FirefoxProfile()
     profile.set_preference("reader.parse-on-load.enabled", False)
     profile.set_preference("pdfjs.disabled", True)
+    if proxy_string:
+        proxy_server = proxy_string.split(':')[0]
+        proxy_port = proxy_string.split(':')[1]
+        profile.set_preference("network.proxy.type", 1)
+        profile.set_preference("network.proxy.http", proxy_server)
+        profile.set_preference("network.proxy.http_port", int(proxy_port))
+        profile.set_preference("network.proxy.ssl", proxy_server)
+        profile.set_preference("network.proxy.ssl_port", int(proxy_port))
     profile.set_preference(
         "security.mixed_content.block_active_content", False)
     profile.set_preference(
@@ -30,15 +40,43 @@ def _create_firefox_profile(downloads_path):
     return profile
 
 
+def display_proxy_warning(proxy_string):
+    message = ('\n\nWARNING: Proxy String ["%s"] is NOT in the expected '
+               '"ip_address:port" format, (OR the key does not exist '
+               'in proxy_list.PROXY_LIST). '
+               '*** DEFAULTING to NOT USING a Proxy Server! ***'
+               % proxy_string)
+    warnings.simplefilter('always', Warning)  # See Warnings
+    warnings.warn(message, category=Warning, stacklevel=2)
+    warnings.simplefilter('default', Warning)  # Set Default
+
+
+def validate_proxy_string(proxy_string):
+    if proxy_string in proxy_list.PROXY_LIST.keys():
+        proxy_string = proxy_list.PROXY_LIST[proxy_string]
+        if not proxy_string:
+            return None
+    valid = re.match('^\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}:\d+$', proxy_string)
+    if valid:
+        proxy_string = valid.group()
+    else:
+        display_proxy_warning(proxy_string)
+        proxy_string = None
+    return proxy_string
+
+
 def get_driver(browser_name, headless=False, use_grid=False,
-               servername='localhost', port=4444):
+               servername='localhost', port=4444, proxy_string=None):
+    if proxy_string:
+        proxy_string = validate_proxy_string(proxy_string)
     if use_grid:
-        return get_remote_driver(browser_name, headless, servername, port)
+        return get_remote_driver(
+            browser_name, headless, servername, port, proxy_string)
     else:
-        return get_local_driver(browser_name, headless)
+        return get_local_driver(browser_name, headless, proxy_string)
 
 
-def get_remote_driver(browser_name, headless, servername, port):
+def get_remote_driver(browser_name, headless, servername, port, proxy_string):
     downloads_path = download_helper.get_downloads_folder()
     download_helper.reset_downloads_folder()
     address = "http://%s:%s/wd/hub" % (servername, port)
@@ -58,6 +96,8 @@ def get_remote_driver(browser_name, headless, servername, port):
         chrome_options.add_argument("--disable-infobars")
         if headless:
             chrome_options.add_argument("--headless")
+        if proxy_string:
+            chrome_options.add_argument('--proxy-server=%s' % proxy_string)
         if settings.START_CHROME_IN_FULL_SCREEN_MODE:
             # Run Chrome in full screen mode on WINDOWS
             chrome_options.add_argument("--start-maximized")
@@ -71,7 +111,7 @@ def get_remote_driver(browser_name, headless, servername, port):
     if browser_name == constants.Browser.FIREFOX:
         try:
             # Use Geckodriver for Firefox if it's on the PATH
-            profile = _create_firefox_profile(downloads_path)
+            profile = _create_firefox_profile(downloads_path, proxy_string)
             firefox_capabilities = DesiredCapabilities.FIREFOX.copy()
             firefox_capabilities['marionette'] = True
             if headless:
@@ -85,7 +125,7 @@ def get_remote_driver(browser_name, headless, servername, port):
                 browser_profile=profile)
         except WebDriverException:
             # Don't use Geckodriver: Only works for old versions of Firefox
-            profile = _create_firefox_profile(downloads_path)
+            profile = _create_firefox_profile(downloads_path, proxy_string)
             firefox_capabilities = DesiredCapabilities.FIREFOX.copy()
             firefox_capabilities['marionette'] = False
             if headless:
@@ -122,7 +162,7 @@ def get_remote_driver(browser_name, headless, servername, port):
                     webdriver.DesiredCapabilities.PHANTOMJS))
 
 
-def get_local_driver(browser_name, headless):
+def get_local_driver(browser_name, headless, proxy_string):
     '''
     Spins up a new web browser and returns the driver.
     Can also be used to spin up additional browsers for the same test.
@@ -134,7 +174,7 @@ def get_local_driver(browser_name, headless):
         try:
             try:
                 # Use Geckodriver for Firefox if it's on the PATH
-                profile = _create_firefox_profile(downloads_path)
+                profile = _create_firefox_profile(downloads_path, proxy_string)
                 firefox_capabilities = DesiredCapabilities.FIREFOX.copy()
                 firefox_capabilities['marionette'] = True
                 options = webdriver.FirefoxOptions()
@@ -145,7 +185,7 @@ def get_local_driver(browser_name, headless):
                     firefox_options=options)
             except WebDriverException:
                 # Don't use Geckodriver: Only works for old versions of Firefox
-                profile = _create_firefox_profile(downloads_path)
+                profile = _create_firefox_profile(downloads_path, proxy_string)
                 firefox_capabilities = DesiredCapabilities.FIREFOX.copy()
                 firefox_capabilities['marionette'] = False
                 firefox_driver = webdriver.Firefox(
@@ -182,12 +222,14 @@ def get_local_driver(browser_name, headless):
             chrome_options.add_argument("--disable-infobars")
             if headless:
                 chrome_options.add_argument("--headless")
+            if proxy_string:
+                chrome_options.add_argument('--proxy-server=%s' % proxy_string)
             if settings.START_CHROME_IN_FULL_SCREEN_MODE:
                 # Run Chrome in full screen mode on WINDOWS
                 chrome_options.add_argument("--start-maximized")
                 # Run Chrome in full screen mode on MAC/Linux
                 chrome_options.add_argument("--kiosk")
-            return webdriver.Chrome(chrome_options=chrome_options)
+            return webdriver.Chrome(options=chrome_options)
         except Exception as e:
             if headless:
                 raise Exception(e)

seleniumbase/fixtures/base_case.py

@@ -1668,6 +1668,7 @@ def setUp(self):
             self.with_page_source = pytest.config.option.with_page_source
             self.servername = pytest.config.option.servername
             self.port = pytest.config.option.port
+            self.proxy_string = pytest.config.option.proxy_string
             self.database_env = pytest.config.option.database_env
             self.log_path = pytest.config.option.log_path
             self.browser = pytest.config.option.browser
@@ -1728,7 +1729,8 @@ def setUp(self):
                                                   self.headless,
                                                   self.use_grid,
                                                   self.servername,
-                                                  self.port)
+                                                  self.port,
+                                                  self.proxy_string)
 
     def __insert_test_result(self, state, err):
         data_payload = TestcaseDataPayload()

seleniumbase/plugins/pytest_plugin.py

@@ -70,13 +70,20 @@ def pytest_addoption(parser):
     parser.addoption('--server', action='store',
                      dest='servername',
                      default='localhost',
-                     help="""Designates the server used by the test.
+                     help="""Designates the Selenium Grid server to use.
                           Default: localhost.""")
     parser.addoption('--port', action='store',
                      dest='port',
                      default='4444',
-                     help="""Designates the port used by the test.
+                     help="""Designates the Selenium Grid port to use.
                           Default: 4444.""")
+    parser.addoption('--proxy', action='store',
+                     dest='proxy_string',
+                     default=None,
+                     help="""Designates the proxy server:port to use.
+                          Format: servername:port.  OR
+                                  A dict key from proxy_list.PROXY_LIST
+                          Default: None.""")
     parser.addoption('--headless', action="store_true",
                      dest='headless',
                      default=False,

seleniumbase/plugins/selenium_plugin.py

@@ -46,13 +46,21 @@ def options(self, parser, env):
         parser.add_option(
             '--server', action='store', dest='servername',
             default='localhost',
-            help="""Designates the server used by the test.
+            help="""Designates the Selenium Grid server to use.
                     Default: localhost.""")
         parser.add_option(
             '--port', action='store', dest='port',
             default='4444',
-            help="""Designates the port used by the test.
+            help="""Designates the Selenium Grid port to use.
                     Default: 4444.""")
+        parser.add_option(
+            '--proxy', action='store',
+            dest='proxy_string',
+            default=None,
+            help="""Designates the proxy server:port to use.
+                    Format: servername:port.  OR
+                            A dict key from proxy_list.PROXY_LIST
+                    Default: None.""")
         parser.add_option(
             '--headless', action="store_true",
             dest='headless',
@@ -102,6 +110,7 @@ def beforeTest(self, test):
         test.test.headless = self.options.headless
         test.test.servername = self.options.servername
         test.test.port = self.options.port
+        test.test.proxy_string = self.options.proxy_string
         test.test.demo_mode = self.options.demo_mode
         test.test.demo_sleep = self.options.demo_sleep
         test.test.highlights = self.options.highlights