[py] simplify driver binary and driver location selecting (#11864)

iampopovich · diemol · web-flow · commit b003857c98cb · 2023-04-05T10:50:50.000+02:00
* simplify driver selection

* lint fixes

* simplify binary selection

* fix tests

* fix test

---------

Co-authored-by: Diego Molina &lt;diemol@users.noreply.github.com&gt;
diff --git a/py/selenium/webdriver/common/selenium_manager.py b/py/selenium/webdriver/common/selenium_manager.py
@@ -41,11 +41,15 @@ def get_binary() -> Path:
 
         :Returns: The Selenium Manager executable location
         """
-        directory = sys.platform
-        if directory == "darwin":
-            directory = "macos"
-        elif directory in ("win32", "cygwin"):
-            directory = "windows"
+        platform = sys.platform
+
+        dirs = {
+            "darwin": "macos",
+            "win32": "windows",
+            "cygwin": "windows",
+        }
+
+        directory = dirs.get(platform) if dirs.get(platform) else platform
 
         file = "selenium-manager.exe" if directory == "windows" else "selenium-manager"
 
@@ -64,12 +68,20 @@ def driver_location(self, browser: str) -> str:
          - browser: which browser to get the driver path for.
         :Returns: The driver path to use
         """
-        allowed = ("chrome", "firefox", "edge", "ie")
-        if browser not in allowed:
-            raise SeleniumManagerException(f"{browser} is not a valid browser.  Choose one of: {allowed}")
 
-        if browser == "ie":
-            browser = "iexplorer"
+        allowed_browsers = {
+            "chrome": "chrome",
+            "firefox": "firefox",
+            "edge": "edge",
+            "ie": "iexplorer",
+        }
+
+        if browser not in allowed_browsers.keys():
+            raise SeleniumManagerException(
+                f"{browser} is not a valid browser.  Choose one of: {list(allowed_browsers.keys())}"
+            )
+
+        browser = allowed_browsers[browser]
 
         binary, browser_flag, browser, output_flag, output = (
             str(self.get_binary()),
diff --git a/py/test/selenium/webdriver/common/selenium_manager_tests.py b/py/test/selenium/webdriver/common/selenium_manager_tests.py
@@ -22,7 +22,7 @@
 
 
 def test_non_supported_browser_raises_sme():
-    msg = r"foo is not a valid browser.  Choose one of: \('chrome', 'firefox', 'edge', 'ie'\)"
+    msg = r"foo is not a valid browser.  Choose one of: \['chrome', 'firefox', 'edge', 'ie'\]"
     with pytest.raises(SeleniumManagerException, match=msg):
         _ = SeleniumManager().driver_location("foo")
 
