use crate::ChromeForTestingManagerError;
use crate::chromedriver::Chromedriver;
use crate::session::Session;
use rootcause::prelude::ResultExt;
use rootcause::{IntoReportCollection, Report, markers::SendSync};
use thirtyfour::prelude::WebDriverError;
use thirtyfour::{ChromeCapabilities, WebDriverBuilder};
#[doc(hidden)]
#[derive(Debug, Clone, Copy)]
pub struct DefaultCaps;
#[doc(hidden)]
#[derive(Debug, Clone, Copy)]
pub struct DefaultConfig;
#[doc(hidden)]
pub struct CapsSetup<F>(F);
#[doc(hidden)]
pub struct ConfigSetup<F>(F);
mod sealed {
pub trait Sealed {}
}
#[doc(hidden)]
pub trait ApplyCaps: sealed::Sealed {
fn apply(self, caps: &mut ChromeCapabilities) -> Result<(), WebDriverError>;
}
impl sealed::Sealed for DefaultCaps {}
impl ApplyCaps for DefaultCaps {
fn apply(self, _caps: &mut ChromeCapabilities) -> Result<(), WebDriverError> {
Ok(())
}
}
impl<F> sealed::Sealed for CapsSetup<F> {}
impl<F> ApplyCaps for CapsSetup<F>
where
F: FnOnce(&mut ChromeCapabilities) -> Result<(), WebDriverError>,
{
fn apply(self, caps: &mut ChromeCapabilities) -> Result<(), WebDriverError> {
self.0(caps)
}
}
#[doc(hidden)]
pub trait ApplyConfig: sealed::Sealed {
fn apply(self, builder: WebDriverBuilder) -> WebDriverBuilder;
}
impl sealed::Sealed for DefaultConfig {}
impl ApplyConfig for DefaultConfig {
fn apply(self, builder: WebDriverBuilder) -> WebDriverBuilder {
builder
}
}
impl<F> sealed::Sealed for ConfigSetup<F> {}
impl<F> ApplyConfig for ConfigSetup<F>
where
F: FnOnce(WebDriverBuilder) -> WebDriverBuilder,
{
fn apply(self, builder: WebDriverBuilder) -> WebDriverBuilder {
self.0(builder)
}
}
pub struct SessionBuilder<'a, C, B> {
chromedriver: &'a Chromedriver,
caps_setup: C,
config_setup: B,
}
impl<'a> SessionBuilder<'a, DefaultCaps, DefaultConfig> {
pub(crate) fn new(chromedriver: &'a Chromedriver) -> Self {
Self {
chromedriver,
caps_setup: DefaultCaps,
config_setup: DefaultConfig,
}
}
}
impl<'a, B> SessionBuilder<'a, DefaultCaps, B> {
pub fn with_caps<F>(self, f: F) -> SessionBuilder<'a, CapsSetup<F>, B>
where
F: FnOnce(&mut ChromeCapabilities) -> Result<(), WebDriverError>,
{
SessionBuilder {
chromedriver: self.chromedriver,
caps_setup: CapsSetup(f),
config_setup: self.config_setup,
}
}
}
impl<'a, C> SessionBuilder<'a, C, DefaultConfig> {
pub fn with_config<F>(self, f: F) -> SessionBuilder<'a, C, ConfigSetup<F>>
where
F: FnOnce(WebDriverBuilder) -> WebDriverBuilder,
{
SessionBuilder {
chromedriver: self.chromedriver,
caps_setup: self.caps_setup,
config_setup: ConfigSetup(f),
}
}
}
impl<C, B> SessionBuilder<'_, C, B>
where
C: ApplyCaps,
B: ApplyConfig,
{
pub async fn run<T, E, F>(self, f: F) -> Result<T, Report<ChromeForTestingManagerError>>
where
F: for<'b> AsyncFnOnce(&'b Session) -> Result<T, E>,
E: IntoReportCollection<SendSync>,
{
use futures::FutureExt;
let chromedriver = self.chromedriver;
let port = chromedriver.port();
let mut caps = chromedriver.mgr.prepare_caps(&chromedriver.loaded)?;
self.caps_setup
.apply(&mut caps)
.context(ChromeForTestingManagerError::ConfigureSessionCapabilities)?;
let builder = thirtyfour::WebDriver::builder(format!("http://localhost:{port}"), caps);
let driver = self
.config_setup
.apply(builder)
.connect()
.await
.context(ChromeForTestingManagerError::StartWebDriverSession { port })?;
let session = Session { driver };
let maybe_panicked = core::panic::AssertUnwindSafe(f(&session))
.catch_unwind()
.await;
let user_result = match maybe_panicked {
Ok(result) => result.context(ChromeForTestingManagerError::RunSessionCallback),
Err(payload) => {
if let Err(quit_err) = session.quit().await {
tracing::error!(
"Failed to quit WebDriver session after user callback panic: {quit_err:?}"
);
}
std::panic::resume_unwind(payload);
}
};
let quit_result = session.quit().await;
match (user_result, quit_result) {
(Ok(value), Ok(())) => Ok(value),
(Ok(_), Err(quit_err)) => Err(quit_err),
(Err(user_err), Ok(())) => Err(user_err),
(Err(mut user_err), Err(quit_err)) => {
tracing::error!(
"Failed to quit WebDriver session after user failure: {quit_err:?}"
);
user_err
.children_mut()
.push(quit_err.into_dynamic().into_cloneable());
Err(user_err)
}
}
}
}