Add the ability to set the User-Agent string for the browser

Author: mdmintz

seleniumbase/core/browser_launcher.py

@@ -83,7 +83,8 @@ def _add_chrome_proxy_extension(
 
 
 def _set_chrome_options(
-        downloads_path, proxy_string, proxy_auth, proxy_user, proxy_pass):
+        downloads_path, proxy_string, proxy_auth,
+        proxy_user, proxy_pass, user_agent):
     chrome_options = webdriver.ChromeOptions()
     prefs = {
         "download.default_directory": downloads_path,
@@ -100,6 +101,8 @@ def _set_chrome_options(
     chrome_options.add_argument("--allow-file-access-from-files")
     chrome_options.add_argument("--allow-insecure-localhost")
     chrome_options.add_argument("--allow-running-insecure-content")
+    if user_agent:
+        chrome_options.add_argument("--user-agent=%s" % user_agent)
     chrome_options.add_argument("--disable-infobars")
     chrome_options.add_argument("--disable-save-password-bubble")
     chrome_options.add_argument("--disable-single-click-autofill")
@@ -113,7 +116,7 @@ def _set_chrome_options(
     return chrome_options
 
 
-def _create_firefox_profile(downloads_path, proxy_string):
+def _create_firefox_profile(downloads_path, proxy_string, user_agent):
     profile = webdriver.FirefoxProfile()
     profile.accept_untrusted_certs = True
     profile.set_preference("reader.parse-on-load.enabled", False)
@@ -128,6 +131,8 @@ def _create_firefox_profile(downloads_path, proxy_string):
         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))
+    if user_agent:
+        profile.set_preference("general.useragent.override", user_agent)
     profile.set_preference(
         "security.mixed_content.block_active_content", False)
     profile.set_preference("security.csp.enable", False)
@@ -195,7 +200,7 @@ def validate_proxy_string(proxy_string):
 
 def get_driver(browser_name, headless=False, use_grid=False,
                servername='localhost', port=4444, proxy_string=None,
-               cap_file=None):
+               user_agent=None, cap_file=None):
     proxy_auth = False
     proxy_user = None
     proxy_pass = None
@@ -224,27 +229,26 @@ def get_driver(browser_name, headless=False, use_grid=False,
     if use_grid:
         return get_remote_driver(
             browser_name, headless, servername, port, proxy_string, proxy_auth,
-            proxy_user, proxy_pass, cap_file)
+            proxy_user, proxy_pass, user_agent, cap_file)
     else:
         return get_local_driver(
             browser_name, headless, proxy_string, proxy_auth,
-            proxy_user, proxy_pass)
+            proxy_user, proxy_pass, user_agent)
 
 
 def get_remote_driver(
         browser_name, headless, servername, port, proxy_string, proxy_auth,
-        proxy_user, proxy_pass, cap_file):
+        proxy_user, proxy_pass, user_agent, cap_file):
     downloads_path = download_helper.get_downloads_folder()
     download_helper.reset_downloads_folder()
     address = "http://%s:%s/wd/hub" % (servername, port)
     desired_caps = {}
     if cap_file:
         desired_caps = capabilities_parser.get_desired_capabilities(cap_file)
-
     if browser_name == constants.Browser.GOOGLE_CHROME:
         chrome_options = _set_chrome_options(
             downloads_path, proxy_string, proxy_auth,
-            proxy_user, proxy_pass)
+            proxy_user, proxy_pass, user_agent)
         if headless:
             if not proxy_auth:
                 # Headless Chrome doesn't support extensions, which are
@@ -264,7 +268,8 @@ def get_remote_driver(
     elif browser_name == constants.Browser.FIREFOX:
         try:
             # Use Geckodriver for Firefox if it's on the PATH
-            profile = _create_firefox_profile(downloads_path, proxy_string)
+            profile = _create_firefox_profile(
+                downloads_path, proxy_string, user_agent)
             firefox_capabilities = DesiredCapabilities.FIREFOX.copy()
             firefox_capabilities['marionette'] = True
             if headless:
@@ -280,7 +285,8 @@ def get_remote_driver(
                 browser_profile=profile)
         except WebDriverException:
             # Don't use Geckodriver: Only works for old versions of Firefox
-            profile = _create_firefox_profile(downloads_path, proxy_string)
+            profile = _create_firefox_profile(
+                downloads_path, proxy_string, user_agent)
             firefox_capabilities = DesiredCapabilities.FIREFOX.copy()
             firefox_capabilities['marionette'] = False
             if headless:
@@ -360,7 +366,7 @@ def get_remote_driver(
 
 def get_local_driver(
         browser_name, headless, proxy_string, proxy_auth,
-        proxy_user, proxy_pass):
+        proxy_user, proxy_pass, user_agent):
     '''
     Spins up a new web browser and returns the driver.
     Can also be used to spin up additional browsers for the same test.
@@ -372,7 +378,8 @@ def get_local_driver(
         try:
             try:
                 # Use Geckodriver for Firefox if it's on the PATH
-                profile = _create_firefox_profile(downloads_path, proxy_string)
+                profile = _create_firefox_profile(
+                    downloads_path, proxy_string, user_agent)
                 firefox_capabilities = DesiredCapabilities.FIREFOX.copy()
                 firefox_capabilities['marionette'] = True
                 options = webdriver.FirefoxOptions()
@@ -392,7 +399,8 @@ def get_local_driver(
                         options=options)
             except WebDriverException:
                 # Don't use Geckodriver: Only works for old versions of Firefox
-                profile = _create_firefox_profile(downloads_path, proxy_string)
+                profile = _create_firefox_profile(
+                    downloads_path, proxy_string, user_agent)
                 firefox_capabilities = DesiredCapabilities.FIREFOX.copy()
                 firefox_capabilities['marionette'] = False
                 firefox_driver = webdriver.Firefox(
@@ -451,7 +459,7 @@ def get_local_driver(
         try:
             chrome_options = _set_chrome_options(
                 downloads_path, proxy_string, proxy_auth,
-                proxy_user, proxy_pass)
+                proxy_user, proxy_pass, user_agent)
             if headless:
                 # Headless Chrome doesn't support extensions, which are
                 # required when using a proxy server that has authentication.

seleniumbase/fixtures/base_case.py

@@ -2348,8 +2348,8 @@ def save_screenshot(self, name, folder=None):
         return page_actions.save_screenshot(self.driver, name, folder)
 
     def get_new_driver(self, browser=None, headless=None,
-                       servername=None, port=None, proxy=None, switch_to=True,
-                       cap_file=None):
+                       servername=None, port=None, proxy=None, agent=None,
+                       switch_to=True, cap_file=None):
         """ This method spins up an extra browser for tests that require
             more than one. The first browser is already provided by tests
             that import base_case.BaseCase from seleniumbase. If parameters
@@ -2399,6 +2399,9 @@ def get_new_driver(self, browser=None, headless=None,
         proxy_string = proxy
         if proxy_string is None:
             proxy_string = self.proxy_string
+        user_agent = agent
+        if user_agent is None:
+            user_agent = self.user_agent
         if cap_file is None:
             cap_file = self.cap_file
         valid_browsers = constants.ValidBrowsers.valid_browsers
@@ -2413,6 +2416,7 @@ def get_new_driver(self, browser=None, headless=None,
                                                  servername=servername,
                                                  port=port,
                                                  proxy_string=proxy_string,
+                                                 user_agent=user_agent,
                                                  cap_file=cap_file)
         self._drivers_list.append(new_driver)
         if switch_to:
@@ -2776,6 +2780,7 @@ def setUp(self):
             self.servername = sb_config.servername
             self.port = sb_config.port
             self.proxy_string = sb_config.proxy_string
+            self.user_agent = sb_config.user_agent
             self.cap_file = sb_config.cap_file
             self.database_env = sb_config.database_env
             self.message_duration = sb_config.message_duration
@@ -2846,6 +2851,7 @@ def setUp(self):
                                           servername=self.servername,
                                           port=self.port,
                                           proxy=self.proxy_string,
+                                          agent=self.user_agent,
                                           switch_to=True,
                                           cap_file=self.cap_file)
         self._default_driver = self.driver

seleniumbase/plugins/pytest_plugin.py

@@ -97,6 +97,12 @@ def pytest_addoption(parser):
                                   username:password@servername:port  OR
                                   A dict key from proxy_list.PROXY_LIST
                           Default: None.""")
+    parser.addoption('--agent', action='store',
+                     dest='user_agent',
+                     default=None,
+                     help="""Designates the User-Agent for the browser to use.
+                          Format: A string.
+                          Default: None.""")
     parser.addoption('--headless', action="store_true",
                      dest='headless',
                      default=False,
@@ -155,6 +161,7 @@ def pytest_configure(config):
     sb_config.data = config.getoption('data')
     sb_config.environment = config.getoption('environment')
     sb_config.with_selenium = config.getoption('with_selenium')
+    sb_config.user_agent = config.getoption('user_agent')
     sb_config.headless = config.getoption('headless')
     sb_config.with_testing_base = config.getoption('with_testing_base')
     sb_config.with_db_reporting = config.getoption('with_db_reporting')

seleniumbase/plugins/selenium_plugin.py

@@ -19,6 +19,7 @@ class SeleniumBrowser(Plugin):
     self.options.server -- the server used by the test (--server)
     self.options.port -- the port used by the test (--port)
     self.options.proxy -- designates the proxy server:port to use. (--proxy)
+    self.options.agent -- designates the User Agent for the browser. (--agent)
     self.options.headless -- the option to run headlessly (--headless)
     self.options.demo_mode -- the option to slow down Selenium (--demo_mode)
     self.options.demo_sleep -- Selenium action delay in DemoMode (--demo_sleep)
@@ -73,6 +74,13 @@ def options(self, parser, env):
                             username:password@servername:port  OR
                             A dict key from proxy_list.PROXY_LIST
                     Default: None.""")
+        parser.add_option(
+            '--agent', action='store',
+            dest='user_agent',
+            default=None,
+            help="""Designates the User-Agent for the browser to use.
+                    Format: A string.
+                    Default: None.""")
         parser.add_option(
             '--headless', action="store_true",
             dest='headless',
@@ -140,6 +148,7 @@ def beforeTest(self, test):
         test.test.servername = self.options.servername
         test.test.port = self.options.port
         test.test.proxy_string = self.options.proxy_string
+        test.test.user_agent = self.options.user_agent
         test.test.demo_mode = self.options.demo_mode
         test.test.demo_sleep = self.options.demo_sleep
         test.test.highlights = self.options.highlights