use std::collections::HashMap;
use std::future::Future;
use std::pin::Pin;
use std::task::{Context, Poll};
use std::time::Duration;
use futures::StreamExt;
use serde::Deserialize;
use tokio::sync::oneshot;
use tokio::time::Sleep;
use zendriver_transport::SessionHandle;
use crate::error::{Result, ZendriverError};
use crate::expect::UrlMatcher;
const DEFAULT_EXPECT_TIMEOUT: Duration = Duration::from_secs(30);
#[derive(Debug, Clone)]
pub struct MatchedRequest {
pub url: String,
pub method: String,
pub headers: HashMap<String, String>,
pub post_data: Option<Vec<u8>>,
pub request_id: String,
}
#[derive(Debug)]
pub struct RequestExpectation {
rx: oneshot::Receiver<MatchedRequest>,
timeout: Duration,
sleep: Option<Pin<Box<Sleep>>>,
}
impl RequestExpectation {
#[must_use]
pub fn timeout(mut self, dur: Duration) -> Self {
self.timeout = dur;
self.sleep = None;
self
}
pub async fn matched(self) -> Result<MatchedRequest> {
self.await
}
}
impl Future for RequestExpectation {
type Output = Result<MatchedRequest>;
fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
match Pin::new(&mut self.rx).poll(cx) {
Poll::Ready(Ok(req)) => return Poll::Ready(Ok(req)),
Poll::Ready(Err(_)) => {
return Poll::Ready(Err(ZendriverError::Timeout(self.timeout)));
}
Poll::Pending => {}
}
let timeout = self.timeout;
let sleep = self
.sleep
.get_or_insert_with(|| Box::pin(tokio::time::sleep(timeout)));
match sleep.as_mut().poll(cx) {
Poll::Ready(()) => Poll::Ready(Err(ZendriverError::Timeout(timeout))),
Poll::Pending => Poll::Pending,
}
}
}
#[derive(Debug, Deserialize)]
struct RequestWillBeSentEvent {
#[serde(rename = "requestId")]
request_id: String,
request: RequestPayload,
}
#[derive(Debug, Deserialize)]
struct RequestPayload {
url: String,
method: String,
#[serde(default)]
headers: HashMap<String, String>,
#[serde(rename = "postData", default)]
post_data: Option<String>,
}
pub(crate) fn register(session: &SessionHandle, matcher: UrlMatcher) -> RequestExpectation {
let (tx, rx) = oneshot::channel();
let mut stream = session.subscribe::<RequestWillBeSentEvent>("Network.requestWillBeSent");
tokio::spawn(async move {
while let Some(ev) = stream.next().await {
if matcher.matches(&ev.request.url) {
let matched = MatchedRequest {
url: ev.request.url,
method: ev.request.method,
headers: ev.request.headers,
post_data: ev.request.post_data.map(String::into_bytes),
request_id: ev.request_id,
};
let _ = tx.send(matched);
return;
}
}
});
RequestExpectation {
rx,
timeout: DEFAULT_EXPECT_TIMEOUT,
sleep: None,
}
}
#[cfg(test)]
#[allow(clippy::panic, clippy::unwrap_used)]
mod tests {
use super::*;
use serde_json::json;
use zendriver_transport::testing::MockConnection;
#[tokio::test]
async fn expect_request_resolves_on_matching_event() {
let (mock, conn) = MockConnection::pair();
let session = SessionHandle::new(conn.clone(), "S1");
let expectation = register(&session, UrlMatcher::from("/api/"));
mock.emit_event_for_session(
"Network.requestWillBeSent",
json!({
"requestId": "R0",
"request": {
"url": "https://example.com/static/app.js",
"method": "GET",
"headers": {},
},
}),
"S1",
)
.await;
mock.emit_event_for_session(
"Network.requestWillBeSent",
json!({
"requestId": "R1",
"request": {
"url": "https://example.com/api/users",
"method": "POST",
"headers": {
"content-type": "application/json",
},
"postData": "{\"name\":\"x\"}",
},
}),
"S1",
)
.await;
let matched = tokio::time::timeout(Duration::from_secs(2), expectation)
.await
.expect("expectation did not resolve within 2s")
.expect("expectation returned Err");
assert_eq!(matched.request_id, "R1");
assert_eq!(matched.url, "https://example.com/api/users");
assert_eq!(matched.method, "POST");
assert_eq!(
matched.headers.get("content-type").map(String::as_str),
Some("application/json"),
);
assert_eq!(
matched.post_data.as_deref(),
Some(b"{\"name\":\"x\"}".as_slice())
);
conn.shutdown();
}
#[tokio::test]
async fn expect_request_times_out() {
let (_mock, conn) = MockConnection::pair();
let session = SessionHandle::new(conn.clone(), "S1");
let expectation =
register(&session, UrlMatcher::from("/api/")).timeout(Duration::from_millis(50));
let res = expectation.await;
match res {
Err(ZendriverError::Timeout(d)) => {
assert_eq!(d, Duration::from_millis(50));
}
other => panic!("expected Timeout(50ms), got {other:?}"),
}
conn.shutdown();
}
}