Refactor capabilities to make everything clearer

Author: stevepryde

thirtyfour/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "thirtyfour"
-version = "0.32.0-rc.1"
+version = "0.32.0-rc.2"
 authors = ["Steve Pryde <steve@stevepryde.com>"]
 edition = "2021"
 license = "MIT OR Apache-2.0"
@@ -33,6 +33,7 @@ futures = "0.3.21"
 http = "0.2.8"
 log = "0.4.17"
 parking_lot = "0.12.1"
+paste = "1.0.9"
 serde = { version = "1.0", features = ["derive"] }
 serde_json = { version = "1.0", features = ["preserve_order"] }
 serde_repr = "0.1.8"

thirtyfour/examples/chrome_options.rs

@@ -13,7 +13,7 @@ async fn main() -> color_eyre::Result<()> {
     color_eyre::install()?;
 
     let mut caps = DesiredCapabilities::chrome();
-    caps.add_chrome_option(
+    caps.insert_browser_option(
         "prefs",
         serde_json::json!({
             "profile.default_content_settings": {

thirtyfour/src/common/capabilities/chrome.rs

@@ -1,11 +1,12 @@
+use std::path::Path;
+
+use paste::paste;
+use serde::Serialize;
+use serde_json::{json, to_value, Value};
+
 use crate::error::WebDriverResult;
-use crate::Capabilities;
 use crate::CapabilitiesHelper;
-use serde::de::DeserializeOwned;
-use serde::Serialize;
-use serde_json::{from_value, json, to_value};
-use std::ops::{Deref, DerefMut};
-use std::path::Path;
+use crate::{BrowserCapabilitiesHelper, Capabilities};
 
 #[derive(Debug, Clone, Serialize)]
 #[serde(transparent)]
@@ -15,113 +16,161 @@ pub struct ChromeCapabilities {
 
 impl Default for ChromeCapabilities {
     fn default() -> Self {
-        let mut capabilities = Capabilities::new();
-        capabilities.insert("browserName".to_string(), json!("chrome"));
-        ChromeCapabilities {
-            capabilities,
+        Self::new()
+    }
+}
+
+macro_rules! chrome_arg_wrapper {
+    ($($fname:ident => $opt:literal),*) => {
+        paste! {
+            $(
+                #[doc = concat!("Set the ", $opt, " option.")]
+                pub fn [<set_ $fname>](&mut self) -> WebDriverResult<()> {
+                    self.add_arg($opt)
+                }
+
+                #[doc = concat!("Unset the ", $opt, " option.")]
+                pub fn [<unset_ $fname>](&mut self) -> WebDriverResult<()> {
+                    self.remove_arg($opt)
+                }
+
+                #[doc = concat!("Return true if the ", $opt, " option is set.")]
+                pub fn [<is_ $fname>](&mut self) -> bool {
+                    self.has_arg($opt)
+                }
+            )*
         }
     }
 }
 
 impl ChromeCapabilities {
     /// Create a new ChromeCapabilities struct.
     pub fn new() -> Self {
-        ChromeCapabilities::default()
-    }
-
-    /// Add the specified Chrome option. This is a helper method for `add_chrome_arg()`.
-    pub fn add_chrome_option<T>(&mut self, key: &str, value: T) -> WebDriverResult<()>
-    where
-        T: Serialize,
-    {
-        self.add_subkey("goog:chromeOptions", key, value)
-    }
-
-    /// Get the specified Chrome option.
-    pub fn get_chrome_option<T>(&self, key: &str) -> T
-    where
-        T: DeserializeOwned + Default,
-    {
-        self.capabilities
-            .get("goog:chromeOptions")
-            .and_then(|options| options.get(key))
-            .and_then(|option| from_value(option.clone()).ok())
-            .unwrap_or_default()
+        let mut capabilities = Capabilities::new();
+        capabilities.insert("browserName".to_string(), json!("chrome"));
+        ChromeCapabilities {
+            capabilities,
+        }
     }
 
     /// Get the current list of command-line arguments to `chromedriver` as a vec.
-    pub fn get_args(&self) -> Vec<String> {
-        self.get_chrome_option("args")
+    pub fn args(&self) -> Vec<String> {
+        self.browser_option("args").unwrap_or_default()
     }
 
     /// Get the current list of Chrome extensions as a vec.
+    ///
     /// Each item is a base64-encoded string containing the .CRX extension file contents.
     /// Use `add_extension()` to add a new extension file.
-    pub fn get_extensions(&self) -> Vec<String> {
-        self.get_chrome_option("extensions")
+    pub fn extensions(&self) -> Vec<String> {
+        self.browser_option("extensions").unwrap_or_default()
     }
 
     /// Get the path to the chrome binary (if one was previously set).
-    pub fn get_binary(&self) -> String {
-        self.get_chrome_option("binary")
+    pub fn binary(&self) -> Option<String> {
+        self.browser_option("binary")
     }
 
     /// Set the path to chrome binary to use.
     pub fn set_binary(&mut self, path: &str) -> WebDriverResult<()> {
-        self.add_chrome_option("binary", path)
+        self.insert_browser_option("binary", path)
+    }
+
+    /// Unset the chrome binary path.
+    pub fn unset_binary(&mut self) {
+        self.remove_browser_option("binary");
     }
 
     /// Get the current debugger address (if one was previously set).
-    pub fn get_debugger_address(&self) -> String {
-        self.get_chrome_option("debuggerAddress")
+    pub fn debugger_address(&self) -> Option<String> {
+        self.browser_option("debuggerAddress")
     }
 
     /// Set the debugger address.
     pub fn set_debugger_address(&mut self, address: &str) -> WebDriverResult<()> {
-        self.add_chrome_option("debuggerAddress", address)
+        self.insert_browser_option("debuggerAddress", address)
+    }
+
+    /// Unset the debugger address.
+    pub fn unset_debugger_address(&mut self) {
+        self.remove_browser_option("debuggerAddress");
     }
 
-    /// Add the specified command-line argument to `chromedriver`. Eg. "--disable-local-storage"
+    /// Add the specified command-line argument to `chromedriver`.
+    ///
+    /// ## Example
+    ///
+    /// ```ignore
+    /// let mut caps = DesiredCapabilities::chrome();
+    /// caps.add_chrome_arg("--disable-local-storage")?;
+    /// ```
+    ///
     /// The full list of switches can be found here:
     /// [https://chromium.googlesource.com/chromium/src/+/master/chrome/common/chrome_switches.cc](https://chromium.googlesource.com/chromium/src/+/master/chrome/common/chrome_switches.cc)
-    pub fn add_chrome_arg(&mut self, arg: &str) -> WebDriverResult<()> {
-        let mut args = self.get_args();
+    pub fn add_arg(&mut self, arg: &str) -> WebDriverResult<()> {
+        let mut args = self.args();
         let arg_string = arg.to_string();
         if !args.contains(&arg_string) {
             args.push(arg_string);
         }
-        self.add_chrome_option("args", to_value(args)?)
+        self.insert_browser_option("args", to_value(args)?)
     }
 
     /// Remove the specified Chrome command-line argument if it had been added previously.
-    pub fn remove_chrome_arg(&mut self, arg: &str) -> WebDriverResult<()> {
-        let mut args = self.get_args();
+    pub fn remove_arg(&mut self, arg: &str) -> WebDriverResult<()> {
+        let mut args = self.args();
         if args.is_empty() {
             Ok(())
         } else {
             args.retain(|v| v != arg);
-            self.add_chrome_option("args", to_value(args)?)
+            self.insert_browser_option("args", to_value(args)?)
         }
     }
 
+    /// Return true if the specified arg is currently set.
+    pub fn has_arg(&self, arg: &str) -> bool {
+        self.args().contains(&arg.to_string())
+    }
+
+    /// Add the specified experimental option.
+    ///
+    /// ## Example
+    /// ```no_run
+    /// use thirtyfour::DesiredCapabilities;
+    /// let mut caps = DesiredCapabilities::chrome();
+    /// caps.add_experimental_option("excludeSwitches", vec!["--enable-logging"]).unwrap();
+    /// ```
+    pub fn add_experimental_option(
+        &mut self,
+        name: impl Into<String>,
+        value: impl Serialize,
+    ) -> WebDriverResult<()> {
+        self.insert_browser_option(name, value)
+    }
+
+    /// Remove the specified experimental option.
+    pub fn remove_experimental_option(&mut self, name: &str) {
+        self.remove_browser_option(name);
+    }
+
     /// Add a base64-encoded extension.
     pub fn add_encoded_extension(&mut self, extension_base64: &str) -> WebDriverResult<()> {
-        let mut extensions = self.get_extensions();
+        let mut extensions = self.extensions();
         let ext_string = extension_base64.to_string();
         if !extensions.contains(&ext_string) {
             extensions.push(ext_string);
         }
-        self.add_chrome_option("extensions", to_value(extensions)?)
+        self.insert_browser_option("extensions", to_value(extensions)?)
     }
 
     /// Remove the specified base64-encoded extension if it had been added previously.
     pub fn remove_encoded_extension(&mut self, extension_base64: &str) -> WebDriverResult<()> {
-        let mut extensions = self.get_extensions();
+        let mut extensions = self.extensions();
         if extensions.is_empty() {
             Ok(())
         } else {
             extensions.retain(|v| v != extension_base64);
-            self.add_chrome_option("extensions", to_value(extensions)?)
+            self.insert_browser_option("extensions", to_value(extensions)?)
         }
     }
 
@@ -140,77 +189,63 @@ impl ChromeCapabilities {
         self.remove_encoded_extension(&b64_contents)
     }
 
-    /// Set the browser to run headless.
-    pub fn set_headless(&mut self) -> WebDriverResult<()> {
-        self.add_chrome_arg("--headless")
-    }
-
-    /// Unset the headless option.
-    pub fn unset_headless(&mut self) -> WebDriverResult<()> {
-        self.remove_chrome_arg("--headless")
-    }
-
-    /// Set disable web security.
-    pub fn set_disable_web_security(&mut self) -> WebDriverResult<()> {
-        self.add_chrome_arg("--disable-web-security")
-    }
-
-    /// Unset disable web security.
-    pub fn unset_disable_web_security(&mut self) -> WebDriverResult<()> {
-        self.remove_chrome_arg("--disable-web-security")
+    /// Get the list of exclude switches.
+    pub fn exclude_switches(&self) -> Vec<String> {
+        self.browser_option("excludeSwitches").unwrap_or_default()
     }
 
-    /// Set ignore certificate errors.
-    pub fn set_ignore_certificate_errors(&mut self) -> WebDriverResult<()> {
-        self.add_chrome_arg("--ignore-certificate-errors")
-    }
-
-    /// Unset ignore certificate errors.
-    pub fn unset_ignore_certificate_errors(&mut self) -> WebDriverResult<()> {
-        self.remove_chrome_arg("--ignore-certificate-errors")
+    /// Add the specified arg to the list of exclude switches.
+    pub fn add_exclude_switch(&mut self, arg: &str) -> WebDriverResult<()> {
+        let mut args = self.exclude_switches();
+        let arg_string = arg.to_string();
+        if !args.contains(&arg_string) {
+            args.push(arg_string);
+        }
+        self.add_experimental_option("excludeSwitches", to_value(args)?)
     }
 
-    pub fn set_no_sandbox(&mut self) -> WebDriverResult<()> {
-        self.add_chrome_arg("--no-sandbox")
+    /// Remove the specified arg from the list of exclude switches.
+    pub fn remove_exclude_switch(&mut self, arg: &str) -> WebDriverResult<()> {
+        let mut args = self.exclude_switches();
+        if args.is_empty() {
+            Ok(())
+        } else {
+            args.retain(|v| v != arg);
+            self.add_experimental_option("excludeSwitches", to_value(args)?)
+        }
     }
 
-    pub fn unset_no_sandbox(&mut self) -> WebDriverResult<()> {
-        self.remove_chrome_arg("--no-sandbox")
+    chrome_arg_wrapper! {
+        headless => "--headless",
+        disable_web_security => "--disable-web-security",
+        ignore_certificate_errors => "--ignore-certificate-errors",
+        no_sandbox => "--no-sandbox",
+        disable_gpu => "--disable-gpu",
+        disable_dev_shm_usage => "--disable-dev-shm-usage",
+        disable_local_storage => "--disable-local-storage"
     }
+}
 
-    pub fn set_disable_gpu(&mut self) -> WebDriverResult<()> {
-        self.add_chrome_arg("--disable-gpu")
+impl CapabilitiesHelper for ChromeCapabilities {
+    fn _get(&self, key: &str) -> Option<&Value> {
+        self.capabilities._get(key)
     }
 
-    pub fn unset_disable_gpu(&mut self) -> WebDriverResult<()> {
-        self.remove_chrome_arg("--disable-gpu")
+    fn _get_mut(&mut self, key: &str) -> Option<&mut Value> {
+        self.capabilities._get_mut(key)
     }
 
-    pub fn set_disable_dev_shm_usage(&mut self) -> WebDriverResult<()> {
-        self.add_chrome_arg("--disable-dev-shm-usage")
+    fn insert_base_capability(&mut self, key: String, value: Value) {
+        self.capabilities.insert_base_capability(key, value);
     }
+}
 
-    pub fn unset_disable_dev_shm_usage(&mut self) -> WebDriverResult<()> {
-        self.remove_chrome_arg("--disable-dev-shm-usage")
-    }
+impl BrowserCapabilitiesHelper for ChromeCapabilities {
+    const KEY: &'static str = "goog:chromeOptions";
 }
 
 impl From<ChromeCapabilities> for Capabilities {
     fn from(caps: ChromeCapabilities) -> Capabilities {
         caps.capabilities
     }
 }
-
-impl Deref for ChromeCapabilities {
-    type Target = Capabilities;
-
-    fn deref(&self) -> &Self::Target {
-        &self.capabilities
-    }
-}
-
-impl DerefMut for ChromeCapabilities {
-    fn deref_mut(&mut self) -> &mut Self::Target {
-        &mut self.capabilities
-    }
-}