wreq/src/response.rs at d366693b059a8ffcb38ad4414198e4035d814009 · 0x676e67/wreq · GitHub

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
use url::Url;

use crate::Body;

#[derive(Debug, Clone, PartialEq)]
pub(crate) struct ResponseUrl(pub Url);

/// Extension trait for http::response::Builder objects
///
/// Allows the user to add a `Url` to the http::Response
pub trait ResponseBuilderExt {
    /// A builder method for the `http::response::Builder` type that allows the user to add a `Url`
    /// to the `http::Response`
    fn url(self, url: Url) -> Self;
}

/// Extension trait for http::Response objects
///
/// Provides methods to extract URL information from HTTP responses
pub trait ResponseExt {
    /// Returns a reference to the `Url` associated with this response, if available.
    fn url(&self) -> Option<&Url>;
}

impl ResponseBuilderExt for http::response::Builder {
    fn url(self, url: Url) -> Self {
        self.extension(ResponseUrl(url))
    }
}

impl ResponseExt for http::Response<Body> {
    fn url(&self) -> Option<&Url> {
        self.extensions().get::<ResponseUrl>().map(|r| &r.0)
    }
}

#[cfg(test)]
mod tests {
    use http::response::Builder;
    use url::Url;

    use super::{ResponseBuilderExt, ResponseExt, ResponseUrl};
    use crate::Body;

    #[test]
    fn test_response_builder_ext() {
        let url = Url::parse("http://example.com").unwrap();
        let response = Builder::new()
            .status(200)
            .url(url.clone())
            .body(())
            .unwrap();

        assert_eq!(
            response.extensions().get::<ResponseUrl>(),
            Some(&ResponseUrl(url))
        );
    }

    #[test]
    fn test_response_ext() {
        let url = Url::parse("http://example.com").unwrap();
        let response = http::Response::builder()
            .status(200)
            .extension(ResponseUrl(url.clone()))
            .body(Body::empty())
            .unwrap();

        assert_eq!(response.url(), Some(&url));
    }
}