spire_core/testing/
site.rs

1use std::collections::{HashMap, HashSet};
2use std::future::{Future, ready};
3use std::sync::{Arc, Mutex};
4
5use ::http::{HeaderMap, StatusCode, Uri};
6use bytes::Bytes;
7
8use super::{MockPage, Visit};
9use crate::backend::{Backend, Session, StatusPolicy};
10use crate::frontier::seed::Seed;
11use crate::http::{Request, Response, ResponseExt};
12use crate::identity::Identity;
13use crate::sync::lock;
14use crate::{Error, Result};
15
16/// Redirects followed at most, as HTTP clients limit them.
17const MAX_REDIRECTS: usize = 10;
18
19/// A [`Backend`] serving fixed [`MockPage`]s by URI, for tests.
20///
21/// URIs without a page get a `404 Not Found`. Statuses are treated as the
22/// real backends treat them, by a [`StatusPolicy`]: by default, an error
23/// status fails the fetch, so the crawler retries or fails the request.
24///
25/// Clones share the record of [`requests`](Self::requests). Crawls with
26/// identities can be tested too, see [`sent`](Self::sent).
27///
28/// # Examples
29///
30/// ```
31/// use spire_core::crawler::Crawler;
32/// use spire_core::frontier::seed::Seed;
33/// use spire_core::testing::{MockPage, MockSite};
34/// use spire_core::worker::from_fn;
35///
36/// # #[tokio::main(flavor = "current_thread")]
37/// # async fn main() -> spire_core::Result<()> {
38/// let site = MockSite::new()
39///     .page("https://example.com/", MockPage::html("<p>Home</p>"))
40///     .page("https://example.com/old", MockPage::redirect("https://example.com/"));
41///
42/// let worker = from_fn(|cx: spire_core::Context<_>| async move { cx.body().await.map(drop) });
43/// let crawler = Crawler::new(site.clone(), worker).with_seeds([
44///     Seed::get("https://example.com/old")?,
45///     Seed::get("https://example.com/missing")?,
46/// ]);
47///
48/// let stats = crawler.run().await?;
49/// assert_eq!((stats.succeeded, stats.failed), (1, 1));
50/// assert_eq!(site.requests().len(), 2);
51/// # Ok(())
52/// # }
53/// ```
54#[derive(Debug, Clone, Default)]
55pub struct MockSite {
56    pages: Arc<HashMap<Uri, MockPage>>,
57    status_policy: StatusPolicy,
58    /// Identities whose proxy refuses every connection.
59    unreachable: Arc<HashSet<String>>,
60    sent: Arc<Mutex<Vec<Sent>>>,
61}
62
63/// A request a [`MockSite`] received, see [`MockSite::sent`].
64#[derive(Debug, Clone)]
65#[non_exhaustive]
66pub struct Sent {
67    /// The requested URI.
68    pub uri: Uri,
69    /// The name of the identity it went out as, if any.
70    pub identity: Option<String>,
71    /// Its headers.
72    pub headers: HeaderMap,
73}
74
75impl MockSite {
76    /// Creates a site without pages.
77    #[must_use]
78    pub fn new() -> Self {
79        Self::default()
80    }
81
82    /// Serves `page` at `uri`, replacing any previous page there.
83    ///
84    /// # Panics
85    ///
86    /// Panics if `uri` is not a valid absolute URI.
87    #[must_use]
88    pub fn page(mut self, uri: &str, page: MockPage) -> Self {
89        Arc::make_mut(&mut self.pages).insert(parse(uri), page);
90        self
91    }
92
93    /// Replaces the [`StatusPolicy`], [`Classify`](StatusPolicy::Classify)
94    /// by default.
95    #[must_use]
96    pub const fn with_status_policy(mut self, policy: StatusPolicy) -> Self {
97        self.status_policy = policy;
98        self
99    }
100
101    /// Fails every fetch through the identity named `name`, as a proxy
102    /// refusing connections would: a retryable error without a status.
103    #[must_use]
104    pub fn unreachable_through(mut self, name: impl Into<String>) -> Self {
105        Arc::make_mut(&mut self.unreachable).insert(name.into());
106        self
107    }
108
109    /// Returns the URIs fetched so far, in order, including clones' fetches.
110    /// Redirects are followed within a fetch, so only the requested URI is
111    /// listed.
112    #[must_use]
113    pub fn requests(&self) -> Vec<Uri> {
114        lock(&self.sent)
115            .iter()
116            .map(|sent| sent.uri.clone())
117            .collect()
118    }
119
120    /// Returns the requests received so far, in order, with their identity
121    /// and headers. Fetches through an
122    /// [unreachable](Self::unreachable_through) identity are included.
123    #[must_use]
124    pub fn sent(&self) -> Vec<Sent> {
125        lock(&self.sent).clone()
126    }
127
128    /// Prepares a [`Visit`] to `seed`, to run one request through a worker.
129    #[must_use]
130    pub fn visit(&self, seed: Seed) -> Visit {
131        Visit::new(self.clone(), seed)
132    }
133
134    /// Serves `uri`, following redirects, and returns the page with the URI
135    /// it was served from.
136    fn serve(&self, uri: &Uri) -> Result<(MockPage, Uri)> {
137        let not_found = || MockPage::status(StatusCode::NOT_FOUND);
138        let mut uri = uri.clone();
139        for _ in 0..=MAX_REDIRECTS {
140            let page = self.pages.get(&uri).cloned().unwrap_or_else(not_found);
141            match page.location {
142                Some(location) => uri = location,
143                None => return Ok((page, uri)),
144            }
145        }
146        Err(Error::new(crate::ErrorKind::Http)
147            .with_message("too many redirects")
148            .with_uri(uri))
149    }
150}
151
152/// Whether `request` is conditional and still matches a page with
153/// `headers`, as a server answering `304 Not Modified` decides: an
154/// `If-None-Match` naming the page's `ETag`, or an `If-Modified-Since` equal
155/// to its `Last-Modified`.
156fn not_modified(request: &Request, headers: &::http::HeaderMap) -> bool {
157    use ::http::header::{ETAG, IF_MODIFIED_SINCE, IF_NONE_MATCH, LAST_MODIFIED};
158
159    let asked = request.headers();
160    if let (Some(etag), Some(none_match)) = (headers.get(ETAG), asked.get(IF_NONE_MATCH)) {
161        let none_match = none_match.to_str().unwrap_or_default();
162        return none_match
163            .split(',')
164            .any(|tag| tag.trim() == "*" || tag.trim().as_bytes() == etag.as_bytes());
165    }
166    matches!(
167        (headers.get(LAST_MODIFIED), asked.get(IF_MODIFIED_SINCE)),
168        (Some(modified), Some(since)) if modified == since
169    )
170}
171
172/// Parses an absolute URI given to a [`MockSite`] or [`MockPage`].
173pub(super) fn parse(uri: &str) -> Uri {
174    match uri.parse::<Uri>() {
175        Ok(parsed) if parsed.scheme().is_some() && parsed.authority().is_some() => {
176            // Normalised as requests are: an empty path becomes `/`.
177            Seed::get(uri).map_or(parsed, |seed| seed.uri().clone())
178        }
179        _ => panic!("`{uri}` is not an absolute URI"),
180    }
181}
182
183impl Backend for MockSite {
184    type Session = MockSession;
185
186    fn connect(&self) -> impl Future<Output = Result<Self::Session>> + Send {
187        ready(Ok(self.session(None)))
188    }
189
190    /// Records each request's identity, see [`MockSite::sent`], and refuses
191    /// those [unreachable](MockSite::unreachable_through). Proxies are
192    /// accepted, as nothing is sent through them.
193    fn connect_as(
194        &self,
195        identity: &Identity,
196    ) -> impl Future<Output = Result<Self::Session>> + Send {
197        ready(Ok(self.session(Some(identity.clone()))))
198    }
199}
200
201impl MockSite {
202    fn session(&self, identity: Option<Identity>) -> MockSession {
203        MockSession {
204            site: self.clone(),
205            identity,
206        }
207    }
208}
209
210/// The [`Session`] of a [`MockSite`], for one request.
211#[derive(Debug)]
212pub struct MockSession {
213    site: MockSite,
214    identity: Option<Identity>,
215}
216
217impl MockSession {
218    async fn serve(&self, request: &Request) -> Result<Response> {
219        let identity = self
220            .identity
221            .as_ref()
222            .map(|identity| identity.name().to_owned());
223        lock(&self.site.sent).push(Sent {
224            uri: request.uri().clone(),
225            identity: identity.clone(),
226            headers: request.headers().clone(),
227        });
228        if identity.is_some_and(|name| self.site.unreachable.contains(&name)) {
229            return Err(Error::new(crate::ErrorKind::Http)
230                .with_message("the proxy refused the connection")
231                .with_uri(request.uri().clone()));
232        }
233        let (page, served_from) = self.site.serve(request.uri())?;
234        if !page.latency.is_zero() {
235            tokio::time::sleep(page.latency).await;
236        }
237
238        let mut head = Response::new(());
239        *head.status_mut() = page.status;
240        let unchanged = page.status.is_success() && not_modified(request, &page.headers);
241        if unchanged {
242            *head.status_mut() = StatusCode::NOT_MODIFIED;
243        }
244        *head.headers_mut() = page.headers;
245        head.set_final_uri(served_from);
246
247        // As the real backends, an error status is rejected with no body.
248        let head = self.site.status_policy.check(head)?;
249        Ok(head.map(|()| if unchanged { Bytes::new() } else { page.body }))
250    }
251}
252
253impl Session for MockSession {
254    fn fetch(&self, request: &Request) -> impl Future<Output = Result<Response>> + Send {
255        self.serve(request)
256    }
257}
spire_core/testing/site.rs

spire_core/testing/
site.rs
