Retry api (#81)

Cargo.toml

@@ -28,7 +28,7 @@ serde = "1.0"
 serde_json = "1.0"
 futures-core = "0.3.0"
 futures-util = "0.3.0"
-tokio = { version = "0.2.0", features = [ "sync" ] }
+tokio = { version = "0.2.0", features = [ "sync", "time" ] }
 hyper = { version = "0.13.0", features = [ "stream" ] }
 hyper-tls = "0.4.0"
 cookie = { version = "0.13", features = ["percent-encode"] }

src/call.rs

@@ -0,0 +1,222 @@
+use crate::{error, Client, Element, Locator};
+
+use futures_util::future::{select, Either};
+use futures_util::pin_mut;
+
+use std::fmt;
+use std::future::Future;
+use std::pin::Pin;
+use std::task::{Context, Poll};
+use std::time::Duration;
+
+use tokio::time::Instant;
+
+use webdriver::command::{LocatorParameters, WebDriverCommand};
+use webdriver::common::WebElement;
+
+type PinBoxFut<T> = Pin<Box<dyn Future<Output = Result<T, error::CmdError>> + Send>>;
+type PinMutFut<'a, T> = Pin<&'a mut (dyn Future<Output = Result<T, error::CmdError>> + Send)>;
+
+mod sealed {
+    use super::PinBoxFut;
+    use crate::{error, Client};
+
+    pub trait Command {
+        type Output;
+        fn invoke(&self, client: Client) -> PinBoxFut<Self::Output>;
+        fn handle_error(error: error::CmdError) -> Result<(), error::CmdError>;
+    }
+}
+
+use sealed::*;
+
+/// TODO
+#[derive(Debug)]
+pub struct FindDescendant {
+    search: LocatorParameters,
+    element: WebElement,
+}
+
+impl Command for FindDescendant {
+    type Output = Element;
+
+    fn invoke(&self, mut client: Client) -> PinBoxFut<Element> {
+        let search = LocatorParameters {
+            using: self.search.using,
+            value: self.search.value.clone(),
+        };
+
+        let cmd = WebDriverCommand::FindElementElement(self.element.clone(), search);
+
+        Box::pin(async move {
+            let res = client.issue(cmd).await?;
+            let e = client.parse_lookup(res)?;
+
+            Ok(Element { client, element: e })
+        })
+    }
+
+    fn handle_error(error: error::CmdError) -> Result<(), error::CmdError> {
+        match error {
+            error::CmdError::NoSuchElement(_) => Ok(()),
+            err => Err(err),
+        }
+    }
+}
+
+/// TODO
+#[derive(Debug)]
+pub struct Find(LocatorParameters);
+
+impl Command for Find {
+    type Output = Element;
+
+    fn invoke(&self, mut client: Client) -> PinBoxFut<Element> {
+        let locator = LocatorParameters {
+            using: self.0.using,
+            value: self.0.value.clone(),
+        };
+
+        Box::pin(async move { client.by(locator).await })
+    }
+
+    fn handle_error(error: error::CmdError) -> Result<(), error::CmdError> {
+        match error {
+            error::CmdError::NoSuchElement(_) => Ok(()),
+            err => Err(err),
+        }
+    }
+}
+
+enum State<T>
+where
+    T: Command,
+{
+    Ready(T),
+    Once(PinBoxFut<T::Output>),
+}
+
+impl<T> fmt::Debug for State<T>
+where
+    T: fmt::Debug + Command,
+{
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        match self {
+            State::Ready(t) => write!(f, "State::Ready({:?})", t),
+            State::Once(_) => write!(f, "State::Once(...)"),
+        }
+    }
+}
+
+impl<T> State<T>
+where
+    T: Command,
+{
+    fn once(&mut self) -> PinMutFut<'_, T::Output> {
+        match self {
+            State::Once(ref mut p) => p.as_mut(),
+            _ => panic!(),
+        }
+    }
+}
+
+/// TODO
+#[derive(Debug)]
+pub struct Retry<T>
+where
+    T: Command,
+{
+    client: Client,
+    state: State<T>,
+}
+
+impl<T> Future for Retry<T>
+where
+    T: Unpin + Command,
+{
+    type Output = Result<T::Output, error::CmdError>;
+
+    fn poll(self: Pin<&mut Self>, ctx: &mut Context<'_>) -> Poll<Self::Output> {
+        let this = self.get_mut();
+        let future = match this.state {
+            State::Ready(ref mut factory) => {
+                this.state = State::Once(factory.invoke(this.client.clone()));
+                this.state.once()
+            }
+            State::Once(ref mut f) => f.as_mut(),
+        };
+
+        future.poll(ctx)
+    }
+}
+
+impl Retry<FindDescendant> {
+    pub(crate) fn find_descendant(
+        client: Client,
+        element: WebElement,
+        search: Locator<'_>,
+    ) -> Self {
+        Self {
+            client,
+            state: State::Ready(FindDescendant {
+                search: search.into(),
+                element,
+            }),
+        }
+    }
+}
+
+impl Retry<Find> {
+    pub(crate) fn find(client: Client, locator: Locator<'_>) -> Self {
+        Self {
+            client,
+            state: State::Ready(Find(locator.into())),
+        }
+    }
+}
+
+impl<T> Retry<T>
+where
+    T: Command,
+{
+    /// TODO
+    pub async fn retry_forever(self) -> Result<T::Output, error::CmdError> {
+        let factory = match self.state {
+            State::Ready(f) => f,
+            _ => panic!(),
+        };
+
+        loop {
+            match factory.invoke(self.client.clone()).await {
+                Ok(x) => return Ok(x),
+                Err(e) => T::handle_error(e)?,
+            }
+        }
+    }
+
+    /// TODO
+    pub async fn retry_for(self, duration: Duration) -> Result<T::Output, error::CmdError> {
+        let a = self.retry_forever();
+        let b = tokio::time::delay_for(duration);
+
+        pin_mut!(a);
+
+        match select(a, b).await {
+            Either::Left(l) => l.0,
+            Either::Right(_) => Err(error::CmdError::RetriesExhausted),
+        }
+    }
+
+    /// TODO
+    pub async fn retry_until(self, deadline: Instant) -> Result<T::Output, error::CmdError> {
+        let a = self.retry_forever();
+        let b = tokio::time::delay_until(deadline);
+
+        pin_mut!(a);
+
+        match select(a, b).await {
+            Either::Left(l) => l.0,
+            Either::Right(_) => Err(error::CmdError::RetriesExhausted),
+        }
+    }
+}

src/error.rs

@@ -117,6 +117,9 @@ pub enum CmdError {
 
     /// Could not decode a base64 image
     ImageDecodeError(::base64::DecodeError),
+
+    /// A timeout or retry limit was reached.
+    RetriesExhausted,
 }
 
 impl CmdError {
@@ -155,6 +158,7 @@ impl Error for CmdError {
             CmdError::NotW3C(..) => "webdriver returned non-conforming response",
             CmdError::InvalidArgument(..) => "invalid argument provided",
             CmdError::ImageDecodeError(..) => "error decoding image",
+            CmdError::RetriesExhausted => "timeout or retry limit reached",
         }
     }
 
@@ -168,7 +172,10 @@ impl Error for CmdError {
             CmdError::Lost(ref e) => Some(e),
             CmdError::Json(ref e) => Some(e),
             CmdError::ImageDecodeError(ref e) => Some(e),
-            CmdError::NotJson(_) | CmdError::NotW3C(_) | CmdError::InvalidArgument(..) => None,
+            CmdError::NotJson(_)
+            | CmdError::NotW3C(_)
+            | CmdError::InvalidArgument(..)
+            | CmdError::RetriesExhausted => None,
         }
     }
 }
@@ -191,6 +198,7 @@ impl fmt::Display for CmdError {
             CmdError::InvalidArgument(ref arg, ref msg) => {
                 write!(f, "Invalid argument `{}`: {}", arg, msg)
             }
+            CmdError::RetriesExhausted => write!(f, "timeout or retry limit reached"),
         }
     }
 }
@@ -242,4 +250,4 @@ mod tests {
         println!("{}", CmdError::NotJson("test".to_string()));
         println!("{}", NewSessionError::Lost(IOError::last_os_error()));
     }
-}
+}

src/lib.rs

@@ -142,13 +142,16 @@ macro_rules! via_json {
 
 pub use hyper::Method;
 
+mod call;
 /// Error types.
 pub mod error;
 
 /// The long-running session future we spawn for multiplexing onto a running WebDriver instance.
 mod session;
 use crate::session::{Cmd, Session};
 
+pub use crate::call::{Find, FindDescendant, Retry};
+
 /// An element locator.
 ///
 /// See <https://www.w3.org/TR/webdriver/#element-retrieval>.
@@ -685,8 +688,8 @@ impl Client {
     }
 
     /// Find an element on the page.
-    pub async fn find(&mut self, search: Locator<'_>) -> Result<Element, error::CmdError> {
-        self.by(search.into()).await
+    pub fn find(&mut self, search: Locator<'_>) -> Retry<Find> {
+        Retry::find(self.clone(), search)
     }
 
     /// Find elements on the page.
@@ -719,29 +722,6 @@ impl Client {
         Ok(())
     }
 
-    /// Wait for the given element to be present on the page.
-    ///
-    /// This can be useful to wait for something to appear on the page before interacting with it.
-    /// While this currently just spins and yields, it may be more efficient than this in the
-    /// future. In particular, in time, it may only run `is_ready` again when an event occurs on
-    /// the page.
-    pub async fn wait_for_find(&mut self, search: Locator<'_>) -> Result<Element, error::CmdError> {
-        let s: webdriver::command::LocatorParameters = search.into();
-        loop {
-            match self
-                .by(webdriver::command::LocatorParameters {
-                    using: s.using,
-                    value: s.value.clone(),
-                })
-                .await
-            {
-                Ok(v) => break Ok(v),
-                Err(error::CmdError::NoSuchElement(_)) => {}
-                Err(e) => break Err(e),
-            }
-        }
-    }
-
     /// Wait for the page to navigate to a new URL before proceeding.
     ///
     /// If the `current` URL is not provided, `self.current_url()` will be used. Note however that
@@ -998,20 +978,10 @@ impl Element {
     }
 
     /// Find the first matching descendant element.
-    pub async fn find(&mut self, search: Locator<'_>) -> Result<Element, error::CmdError> {
-        let res = self
-            .client
-            .issue(WebDriverCommand::FindElementElement(
-                self.element.clone(),
-                search.into(),
-            ))
-            .await?;
-        let e = self.client.parse_lookup(res)?;
-        Ok(Element {
-            client: self.client.clone(),
-            element: e,
-        })
+    pub fn find(&mut self, search: Locator<'_>) -> Retry<FindDescendant> {
+        Retry::find_descendant(self.client.clone(), self.element.clone(), search)
     }
+
     /// Find all matching descendant elements.
     pub async fn find_all(&mut self, search: Locator<'_>) -> Result<Vec<Element>, error::CmdError> {
         let res = self
@@ -1171,7 +1141,7 @@ impl Form {
 
     /// Find a form input with the given `name` and set its value to `value`.
     pub async fn set_by_name(&mut self, field: &str, value: &str) -> Result<Self, error::CmdError> {
-        let locator = format!("input[name='{}']", field);
+        let locator = format!("[name='{}']", field);
         let locator = Locator::Css(&locator);
         self.set(locator, value).await
     }