use tokio::time::{delay_for, Duration, Instant};
use futures::Future;
use serde::{Deserialize, Serialize};
use std::mem;
use std::pin::Pin;
use std::sync::Arc;
use stringmatch::Needle;
use thirtyfour::error::{WebDriverError, WebDriverErrorInfo};
use thirtyfour::prelude::{WebDriver, WebDriverResult};
use thirtyfour::{By, WebDriverCommands, WebDriverSession, WebElement};
fn get_selector_summary(selectors: &[ElementSelector]) -> String {
let criteria: Vec<String> = selectors.iter().map(|s| s.by.to_string()).collect();
format!("[{}]", criteria.join(","))
}
fn no_such_element(selectors: &[ElementSelector]) -> WebDriverError {
WebDriverError::NoSuchElement(WebDriverErrorInfo::new(&format!(
"Element(s) not found using selectors: {}",
&get_selector_summary(selectors)
)))
}
#[derive(Debug, Serialize, Deserialize, Clone)]
pub enum ElementPoller {
NoWait,
TimeoutWithInterval(Duration, Duration),
NumTriesWithInterval(u32, Duration),
TimeoutWithIntervalAndMinTries(Duration, Duration, u32),
}
pub(crate) type ElementPredicate = Box<
dyn for<'a> Fn(
&'a WebElement<'a>,
) -> Pin<Box<dyn Future<Output = WebDriverResult<bool>> + Send + 'a>>
+ Send
+ Sync
+ 'static,
>;
pub struct ElementSelector<'a> {
pub single: bool,
pub by: By<'a>,
pub filters: Vec<ElementPredicate>,
}
impl<'a> ElementSelector<'a> {
pub fn new(by: By<'a>) -> Self {
Self {
single: false,
by: by.clone(),
filters: Vec::new(),
}
}
pub fn set_single(&mut self) {
self.single = true;
}
pub fn add_filter(&mut self, f: ElementPredicate) {
self.filters.push(f);
}
pub async fn run_filters<'b>(
&self,
mut elements: Vec<WebElement<'b>>,
) -> WebDriverResult<Vec<WebElement<'b>>> {
for func in &self.filters {
let tmp_elements = mem::replace(&mut elements, Vec::new());
for element in tmp_elements {
if func(&element).await? {
elements.push(element);
}
}
if elements.is_empty() {
break;
}
}
Ok(elements)
}
}
pub enum ElementQuerySource<'a> {
Driver(&'a WebDriverSession),
Element(&'a WebElement<'a>),
}
pub struct ElementQuery<'a> {
source: Arc<ElementQuerySource<'a>>,
poller: ElementPoller,
selectors: Vec<ElementSelector<'a>>,
}
impl<'a> ElementQuery<'a> {
fn new(source: ElementQuerySource<'a>, poller: ElementPoller, by: By<'a>) -> Self {
let selector = ElementSelector::new(by.clone());
Self {
source: Arc::new(source),
poller,
selectors: vec![selector],
}
}
pub fn with_poller(mut self, poller: ElementPoller) -> Self {
self.poller = poller;
self
}
pub fn wait(self, timeout: Duration, interval: Duration) -> Self {
self.with_poller(ElementPoller::TimeoutWithInterval(timeout, interval))
}
pub fn nowait(self) -> Self {
self.with_poller(ElementPoller::NoWait)
}
fn add_selector(mut self, selector: ElementSelector<'a>) -> Self {
self.selectors.push(selector);
self
}
pub fn or(self, by: By<'a>) -> Self {
self.add_selector(ElementSelector::new(by))
}
pub async fn exists(&self) -> WebDriverResult<bool> {
let elements = self.run_poller(false).await?;
Ok(!elements.is_empty())
}
pub async fn not_exists(&self) -> WebDriverResult<bool> {
let elements = self.run_poller(true).await?;
Ok(elements.is_empty())
}
pub async fn first(&self) -> WebDriverResult<WebElement<'a>> {
let mut elements = self.run_poller(false).await?;
if elements.is_empty() {
Err(no_such_element(&self.selectors))
} else {
Ok(elements.remove(0))
}
}
pub async fn all(&self) -> WebDriverResult<Vec<WebElement<'a>>> {
self.run_poller(false).await
}
pub async fn all_required(&self) -> WebDriverResult<Vec<WebElement<'a>>> {
let elements = self.run_poller(false).await?;
if elements.is_empty() {
Err(no_such_element(&self.selectors))
} else {
Ok(elements)
}
}
async fn run_poller(&self, inverted: bool) -> WebDriverResult<Vec<WebElement<'a>>> {
match self.poller {
ElementPoller::NoWait => self.run_poller_with_options(None, None, 0, inverted).await,
ElementPoller::TimeoutWithInterval(timeout, interval) => {
self.run_poller_with_options(Some(timeout), Some(interval), 0, inverted).await
}
ElementPoller::NumTriesWithInterval(max_tries, interval) => {
self.run_poller_with_options(None, Some(interval), max_tries, inverted).await
}
ElementPoller::TimeoutWithIntervalAndMinTries(timeout, interval, min_tries) => {
self.run_poller_with_options(Some(timeout), Some(interval), min_tries, inverted)
.await
}
}
}
async fn run_poller_with_options(
&self,
timeout: Option<Duration>,
interval: Option<Duration>,
min_tries: u32,
inverted: bool,
) -> WebDriverResult<Vec<WebElement<'a>>> {
let no_such_element_error = no_such_element(&self.selectors);
if self.selectors.is_empty() {
return Err(no_such_element_error);
}
let mut tries = 0;
let check = |value: bool| {
if inverted {
!value
} else {
value
}
};
let start = Instant::now();
loop {
tries += 1;
for selector in &self.selectors {
let mut elements = match self.fetch_elements_from_source(selector).await {
Ok(x) => x,
Err(WebDriverError::NoSuchElement(_)) => Vec::new(),
Err(e) => return Err(e),
};
if !elements.is_empty() {
elements = selector.run_filters(elements).await?;
}
if check(!elements.is_empty()) {
return Ok(elements);
}
if let Some(t) = timeout {
if start.elapsed() >= t && tries >= min_tries {
return Ok(Vec::new());
}
}
}
if timeout.is_none() && tries >= min_tries {
return Ok(Vec::new());
}
if let Some(i) = interval {
let minimum_elapsed = i * tries;
let actual_elapsed = start.elapsed();
if actual_elapsed < minimum_elapsed {
delay_for(minimum_elapsed - actual_elapsed).await;
}
}
}
}
fn fetch_elements_from_source(
&self,
selector: &ElementSelector<'a>,
) -> impl Future<Output = WebDriverResult<Vec<WebElement<'a>>>> + Send {
let by = selector.by.clone();
let single = selector.single;
let source = self.source.clone();
async move {
match single {
true => match source.as_ref() {
ElementQuerySource::Driver(driver) => {
driver.find_element(by).await.map(|x| vec![x])
}
ElementQuerySource::Element(element) => {
element.find_element(by).await.map(|x| vec![x])
}
},
false => match source.as_ref() {
ElementQuerySource::Driver(driver) => driver.find_elements(by).await,
ElementQuerySource::Element(element) => element.find_elements(by).await,
},
}
}
}
pub fn with_filter(mut self, f: ElementPredicate) -> Self {
if let Some(selector) = self.selectors.last_mut() {
selector.add_filter(f);
}
self
}
pub fn with_single_selector(mut self) -> Self {
if let Some(selector) = self.selectors.last_mut() {
selector.set_single();
}
self
}
pub fn and_enabled(self) -> Self {
self.with_filter(Box::new(|elem| {
Box::pin(async move { elem.is_enabled().await.or(Ok(false)) })
}))
}
pub fn and_not_enabled(self) -> Self {
self.with_filter(Box::new(|elem| {
Box::pin(async move { elem.is_enabled().await.map(|x| !x).or(Ok(false)) })
}))
}
pub fn and_selected(self) -> Self {
self.with_filter(Box::new(|elem| {
Box::pin(async move { elem.is_selected().await.or(Ok(false)) })
}))
}
pub fn and_not_selected(self) -> Self {
self.with_filter(Box::new(|elem| {
Box::pin(async move { elem.is_selected().await.map(|x| !x).or(Ok(false)) })
}))
}
pub fn and_displayed(self) -> Self {
self.with_filter(Box::new(|elem| {
Box::pin(async move { elem.is_displayed().await.or(Ok(false)) })
}))
}
pub fn and_not_displayed(self) -> Self {
self.with_filter(Box::new(|elem| {
Box::pin(async move { elem.is_displayed().await.map(|x| !x).or(Ok(false)) })
}))
}
pub fn and_clickable(self) -> Self {
self.with_filter(Box::new(|elem| {
Box::pin(async move { elem.is_clickable().await.or(Ok(false)) })
}))
}
pub fn and_not_clickable(self) -> Self {
self.with_filter(Box::new(|elem| {
Box::pin(async move { elem.is_clickable().await.map(|x| !x).or(Ok(false)) })
}))
}
pub fn with_text<N>(self, text: N) -> Self
where
N: Needle + Clone + Send + Sync + 'static,
{
self.with_filter(Box::new(move |elem| {
let text = text.clone();
Box::pin(async move { elem.text().await.map(|x| text.is_match(&x)).or(Ok(false)) })
}))
}
pub fn with_id<N>(self, id: N) -> Self
where
N: Needle + Clone + Send + Sync + 'static,
{
self.with_filter(Box::new(move |elem| {
let id = id.clone();
Box::pin(async move {
match elem.id().await {
Ok(Some(x)) => Ok(id.is_match(&x)),
_ => Ok(false),
}
})
}))
}
pub fn with_class<N>(self, class_name: N) -> Self
where
N: Needle + Clone + Send + Sync + 'static,
{
self.with_filter(Box::new(move |elem| {
let class_name = class_name.clone();
Box::pin(async move {
match elem.class_name().await {
Ok(Some(x)) => Ok(class_name.is_match(&x)),
_ => Ok(false),
}
})
}))
}
pub fn with_tag<N>(self, tag_name: N) -> Self
where
N: Needle + Clone + Send + Sync + 'static,
{
self.with_filter(Box::new(move |elem| {
let tag_name = tag_name.clone();
Box::pin(
async move { elem.tag_name().await.map(|x| tag_name.is_match(&x)).or(Ok(false)) },
)
}))
}
pub fn with_value<N>(self, value: N) -> Self
where
N: Needle + Clone + Send + Sync + 'static,
{
self.with_filter(Box::new(move |elem| {
let value = value.clone();
Box::pin(async move {
match elem.value().await {
Ok(Some(x)) => Ok(value.is_match(&x)),
_ => Ok(false),
}
})
}))
}
pub fn with_attribute<N>(self, attribute_name: &str, value: N) -> Self
where
N: Needle + Clone + Send + Sync + 'static,
{
let attribute_name = attribute_name.to_string();
self.with_filter(Box::new(move |elem| {
let attribute_name = attribute_name.clone();
let value = value.clone();
Box::pin(async move {
match elem.get_attribute(&attribute_name).await {
Ok(Some(x)) => Ok(value.is_match(&x)),
_ => Ok(false),
}
})
}))
}
pub fn with_attributes<N>(self, desired_attributes: &'static [(String, N)]) -> Self
where
N: Needle + Clone + Send + Sync + 'static,
{
self.with_filter(Box::new(move |elem| {
Box::pin(async move {
for (attribute_name, value) in desired_attributes {
match elem.get_attribute(&attribute_name).await {
Ok(Some(x)) => {
if !value.is_match(&x) {
return Ok(false);
}
}
_ => return Ok(false),
}
}
Ok(true)
})
}))
}
pub fn with_property<N>(self, property_name: &str, value: N) -> Self
where
N: Needle + Clone + Send + Sync + 'static,
{
let property_name = property_name.to_string();
self.with_filter(Box::new(move |elem| {
let property_name = property_name.clone();
let value = value.clone();
Box::pin(async move {
match elem.get_property(&property_name).await {
Ok(Some(x)) => Ok(value.is_match(&x)),
_ => Ok(false),
}
})
}))
}
pub fn with_properties<N>(self, desired_properties: &'static [(&str, N)]) -> Self
where
N: Needle + Clone + Send + Sync + 'static,
{
self.with_filter(Box::new(move |elem| {
Box::pin(async move {
for (property_name, value) in desired_properties {
match elem.get_property(property_name).await {
Ok(Some(x)) => {
if !value.is_match(&x) {
return Ok(false);
}
}
_ => return Ok(false),
}
}
Ok(true)
})
}))
}
pub fn with_css_property<N>(self, css_property_name: &str, value: N) -> Self
where
N: Needle + Clone + Send + Sync + 'static,
{
let css_property_name = css_property_name.to_string();
self.with_filter(Box::new(move |elem| {
let css_property_name = css_property_name.clone();
let value = value.clone();
Box::pin(async move {
match elem.get_css_property(&css_property_name).await {
Ok(x) => Ok(value.is_match(&x)),
_ => Ok(false),
}
})
}))
}
pub fn with_css_properties<N>(self, desired_css_properties: &'static [(&str, N)]) -> Self
where
N: Needle + Clone + Send + Sync + 'static,
{
self.with_filter(Box::new(move |elem| {
Box::pin(async move {
for (css_property_name, value) in desired_css_properties {
match elem.get_css_property(css_property_name).await {
Ok(x) => {
if !value.is_match(&x) {
return Ok(false);
}
}
_ => return Ok(false),
}
}
Ok(true)
})
}))
}
}
pub trait ElementQueryable {
fn query<'a>(&'a self, by: By<'a>) -> ElementQuery<'a>;
}
impl ElementQueryable for WebElement<'_> {
fn query<'a>(&'a self, by: By<'a>) -> ElementQuery<'a> {
let poller: ElementPoller =
self.session.config().get("ElementPoller").unwrap_or(ElementPoller::NoWait);
ElementQuery::new(ElementQuerySource::Element(&self), poller, by)
}
}
impl ElementQueryable for WebDriver {
fn query<'a>(&'a self, by: By<'a>) -> ElementQuery<'a> {
let poller: ElementPoller =
self.config().get("ElementPoller").unwrap_or(ElementPoller::NoWait);
ElementQuery::new(ElementQuerySource::Driver(&self.session), poller, by)
}
}
#[cfg(test)]
async fn _test_is_send() -> WebDriverResult<()> {
use thirtyfour::prelude::*;
fn is_send<T: Send>() {}
fn is_send_val<T: Send>(_val: &T) {}
let selector = ElementSelector::new(By::Css("div"));
is_send_val(&selector.run_filters(Vec::new()));
let caps = DesiredCapabilities::chrome();
let driver = WebDriver::new("http://localhost:4444", &caps).await?;
let query = driver.query(By::Css("div"));
is_send_val(&query.exists());
is_send_val(&query.not_exists());
is_send_val(&query.first());
is_send_val(&query.all());
is_send_val(&query.all_required());
Ok(())
}