spire_core/testing/
site.rs

1use std::collections::HashMap;
2use std::future::{Future, ready};
3use std::sync::{Arc, Mutex};
4
5use bytes::Bytes;
6use http::{StatusCode, Uri};
7
8use super::{MockPage, Visit};
9use crate::backend::{Backend, Client, StatusPolicy};
10use crate::frontier::seed::Seed;
11use crate::request::{Request, Response, ResponseExt};
12use crate::sync::lock;
13use crate::{Error, Result};
14
15/// Redirects followed at most, as HTTP clients limit them.
16const MAX_REDIRECTS: usize = 10;
17
18/// A [`Backend`] serving fixed [`MockPage`]s by URI, for tests.
19///
20/// URIs without a page get a `404 Not Found`. Statuses are treated as the
21/// real backends treat them, by a [`StatusPolicy`]: by default, an error
22/// status fails the fetch, so the crawler retries or fails the request.
23///
24/// Clones share the record of [`requests`](Self::requests).
25///
26/// # Examples
27///
28/// ```
29/// use spire_core::Crawler;
30/// use spire_core::frontier::seed::Seed;
31/// use spire_core::testing::{MockPage, MockSite};
32/// use spire_core::worker::from_fn;
33///
34/// # #[tokio::main(flavor = "current_thread")]
35/// # async fn main() -> spire_core::Result<()> {
36/// let site = MockSite::new()
37///     .page("https://example.com/", MockPage::html("<p>Home</p>"))
38///     .page("https://example.com/old", MockPage::redirect("https://example.com/"));
39///
40/// let worker = from_fn(|cx: spire_core::Context<_>| async move { cx.body().await.map(drop) });
41/// let crawler = Crawler::new(site.clone(), worker).with_seeds([
42///     Seed::get("https://example.com/old")?,
43///     Seed::get("https://example.com/missing")?,
44/// ]);
45///
46/// let stats = crawler.run().await?;
47/// assert_eq!((stats.succeeded, stats.failed), (1, 1));
48/// assert_eq!(site.requests().len(), 2);
49/// # Ok(())
50/// # }
51/// ```
52#[derive(Debug, Clone, Default)]
53pub struct MockSite {
54    pages: Arc<HashMap<Uri, MockPage>>,
55    status_policy: StatusPolicy,
56    requests: Arc<Mutex<Vec<Uri>>>,
57}
58
59impl MockSite {
60    /// Creates a site without pages.
61    #[must_use]
62    pub fn new() -> Self {
63        Self::default()
64    }
65
66    /// Serves `page` at `uri`, replacing any previous page there.
67    ///
68    /// # Panics
69    ///
70    /// Panics if `uri` is not a valid absolute URI.
71    #[must_use]
72    pub fn page(mut self, uri: &str, page: MockPage) -> Self {
73        Arc::make_mut(&mut self.pages).insert(parse(uri), page);
74        self
75    }
76
77    /// Replaces the [`StatusPolicy`], [`Classify`](StatusPolicy::Classify)
78    /// by default.
79    #[must_use]
80    pub const fn with_status_policy(mut self, policy: StatusPolicy) -> Self {
81        self.status_policy = policy;
82        self
83    }
84
85    /// Returns the URIs fetched so far, in order, including clones' fetches.
86    /// Redirects are followed within a fetch, so only the requested URI is
87    /// listed.
88    #[must_use]
89    pub fn requests(&self) -> Vec<Uri> {
90        lock(&self.requests).clone()
91    }
92
93    /// Prepares a [`Visit`] to `seed`, to run one request through a worker.
94    #[must_use]
95    pub fn visit(&self, seed: Seed) -> Visit {
96        Visit::new(self.clone(), seed)
97    }
98
99    /// Serves `uri`, following redirects, and returns the page with the URI
100    /// it was served from.
101    fn serve(&self, uri: &Uri) -> Result<(MockPage, Uri)> {
102        let not_found = || MockPage::status(StatusCode::NOT_FOUND);
103        let mut uri = uri.clone();
104        for _ in 0..=MAX_REDIRECTS {
105            let page = self.pages.get(&uri).cloned().unwrap_or_else(not_found);
106            match page.location {
107                Some(location) => uri = location,
108                None => return Ok((page, uri)),
109            }
110        }
111        Err(Error::new(crate::ErrorKind::Http)
112            .with_message("too many redirects")
113            .with_uri(uri))
114    }
115}
116
117/// Parses an absolute URI given to a [`MockSite`] or [`MockPage`].
118pub(super) fn parse(uri: &str) -> Uri {
119    match uri.parse::<Uri>() {
120        Ok(parsed) if parsed.scheme().is_some() && parsed.authority().is_some() => {
121            // Normalised as requests are: an empty path becomes `/`.
122            Seed::get(uri).map_or(parsed, |seed| seed.uri().clone())
123        }
124        _ => panic!("`{uri}` is not an absolute URI"),
125    }
126}
127
128impl Backend for MockSite {
129    type Client = MockClient;
130
131    fn connect(&self) -> impl Future<Output = Result<Self::Client>> + Send {
132        ready(Ok(MockClient {
133            site: self.clone(),
134            body: Mutex::new(None),
135        }))
136    }
137}
138
139/// The [`Client`] of a [`MockSite`], for one request.
140#[derive(Debug)]
141pub struct MockClient {
142    site: MockSite,
143    body: Mutex<Option<Bytes>>,
144}
145
146impl MockClient {
147    async fn serve(&self, request: &Request) -> Result<Response<()>> {
148        lock(&self.site.requests).push(request.uri().clone());
149        let (page, served_from) = self.site.serve(request.uri())?;
150        if !page.latency.is_zero() {
151            tokio::time::sleep(page.latency).await;
152        }
153
154        let mut head = Response::new(());
155        *head.status_mut() = page.status;
156        *head.headers_mut() = page.headers;
157        head.set_final_uri(served_from);
158
159        // As the real backends, an error status is rejected with no body.
160        let head = self.site.status_policy.check(head)?;
161        *lock(&self.body) = Some(page.body);
162        Ok(head)
163    }
164}
165
166impl Client for MockClient {
167    fn fetch(&self, request: &Request) -> impl Future<Output = Result<Response<()>>> + Send {
168        self.serve(request)
169    }
170
171    fn body(&self) -> impl Future<Output = Result<Bytes>> + Send {
172        let body = lock(&self.body).clone();
173        ready(body.ok_or_else(|| Error::worker("the request was not fetched yet")))
174    }
175}
spire_core/testing/site.rs

spire_core/testing/
site.rs
