kumo 0.2.10

An async web crawling framework for Rust - Scrapy for Rust
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95

use std::collections::HashMap;

use async_trait::async_trait;
use reqwest::header::HeaderMap;

use super::Fetcher;
use crate::{
    error::KumoError,
    extract::{Response, response::ResponseBody},
    middleware::FetchRequest,
};

/// A test-only fetcher that returns predefined responses by URL.
///
/// Use with [`CrawlEngine::fetcher`](crate::engine::CrawlEngine::fetcher) to run
/// spiders against fixed HTML without any network access.
///
/// # Example
/// ```rust,ignore
/// let mock = MockFetcher::new()
///     .with_response("https://example.com", 200, "<h1>Hello</h1>")
///     .with_response("https://example.com/page/2", 200, "<h1>Page 2</h1>");
///
/// let stats = CrawlEngine::builder()
///     .fetcher(mock)
///     .run(MySpider)
///     .await?;
/// ```
#[derive(Debug)]
pub struct MockFetcher {
    responses: HashMap<String, (u16, String)>,
    default: Option<(u16, String)>,
}

impl MockFetcher {
    /// Create an empty mock fetcher. Unregistered URLs return 404 with an empty body.
    pub fn new() -> Self {
        Self {
            responses: HashMap::new(),
            default: None,
        }
    }

    /// Register a response for a specific URL.
    pub fn with_response(mut self, url: &str, status: u16, body: impl Into<String>) -> Self {
        self.responses
            .insert(url.to_string(), (status, body.into()));
        self
    }

    /// Load a response body from a local HTML file for a specific URL.
    /// Panics if the file cannot be read.
    pub fn with_html_file(mut self, url: &str, path: impl AsRef<std::path::Path>) -> Self {
        let body = std::fs::read_to_string(path.as_ref()).unwrap_or_else(|e| {
            panic!(
                "MockFetcher: failed to read {}: {e}",
                path.as_ref().display()
            )
        });
        self.responses.insert(url.to_string(), (200, body));
        self
    }

    /// Respond with this status + body for any URL not explicitly registered.
    pub fn with_default(mut self, status: u16, body: impl Into<String>) -> Self {
        self.default = Some((status, body.into()));
        self
    }
}

impl Default for MockFetcher {
    fn default() -> Self {
        Self::new()
    }
}

#[async_trait]
impl Fetcher for MockFetcher {
    async fn fetch(&self, request: &FetchRequest) -> Result<Response, KumoError> {
        let (status, body) = self
            .responses
            .get(request.url())
            .or(self.default.as_ref())
            .cloned()
            .unwrap_or((404, String::new()));

        Ok(Response::new(
            request.url().to_string(),
            status,
            HeaderMap::new(),
            std::time::Duration::ZERO,
            ResponseBody::Text(body),
        ))
    }
}