Skip to main content

rama_http/layer/classify/
status_in_range_is_error.rs

1use super::{ClassifiedResponse, ClassifyResponse, NeverClassifyEos, SharedClassifier};
2use rama_http_types::StatusCode;
3use std::{fmt, ops::RangeInclusive};
4
5/// Response classifier that considers responses with a status code within some range to be
6/// failures.
7#[derive(Debug, Clone)]
8pub struct StatusInRangeAsFailures {
9    range: RangeInclusive<u16>,
10}
11
12impl StatusInRangeAsFailures {
13    /// Creates a new `StatusInRangeAsFailures`.
14    ///
15    /// # Panics
16    ///
17    /// Panics if the start or end of `range` aren't valid status codes as determined by
18    /// [`StatusCode::from_u16`].
19    ///
20    /// [`StatusCode::from_u16`]: https://docs.rs/http/latest/http/status/struct.StatusCode.html#method.from_u16
21    #[must_use]
22    pub fn new(range: RangeInclusive<u16>) -> Self {
23        assert!(
24            StatusCode::from_u16(*range.start()).is_ok(),
25            "range start isn't a valid status code"
26        );
27        assert!(
28            StatusCode::from_u16(*range.end()).is_ok(),
29            "range end isn't a valid status code"
30        );
31
32        Self { range }
33    }
34
35    /// Creates a new `StatusInRangeAsFailures` that classifies client and server responses as
36    /// failures.
37    ///
38    /// This is a convenience for `StatusInRangeAsFailures::new(400..=599)`.
39    #[must_use]
40    pub fn new_for_client_and_server_errors() -> Self {
41        Self::new(400..=599)
42    }
43
44    /// Convert this `StatusInRangeAsFailures` into a [`MakeClassifier`].
45    ///
46    /// [`MakeClassifier`]: super::MakeClassifier
47    #[must_use]
48    pub fn into_make_classifier(self) -> SharedClassifier<Self> {
49        SharedClassifier::new(self)
50    }
51}
52
53impl ClassifyResponse for StatusInRangeAsFailures {
54    type FailureClass = StatusInRangeFailureClass;
55    type ClassifyEos = NeverClassifyEos<Self::FailureClass>;
56
57    fn classify_response<B>(
58        self,
59        res: &rama_http_types::Response<B>,
60    ) -> ClassifiedResponse<Self::FailureClass, Self::ClassifyEos> {
61        if self.range.contains(&res.status().as_u16()) {
62            let class = StatusInRangeFailureClass::StatusCode(res.status());
63            ClassifiedResponse::Ready(Err(class))
64        } else {
65            ClassifiedResponse::Ready(Ok(()))
66        }
67    }
68
69    fn classify_error<E>(self, error: &E) -> Self::FailureClass
70    where
71        E: std::fmt::Display,
72    {
73        StatusInRangeFailureClass::Error(error.to_string())
74    }
75}
76
77/// The failure class for [`StatusInRangeAsFailures`].
78#[derive(Debug)]
79pub enum StatusInRangeFailureClass {
80    /// A response was classified as a failure with the corresponding status.
81    StatusCode(StatusCode),
82    /// A response was classified as an error with the corresponding error description.
83    Error(String),
84}
85
86impl fmt::Display for StatusInRangeFailureClass {
87    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
88        match self {
89            Self::StatusCode(code) => write!(f, "Status code: {code}"),
90            Self::Error(error) => write!(f, "Error: {error}"),
91        }
92    }
93}
94
95#[cfg(test)]
96mod tests {
97    use super::*;
98    use rama_http_types::Response;
99
100    #[test]
101    fn basic() {
102        let classifier = StatusInRangeAsFailures::new(400..=599);
103
104        assert!(matches!(
105            classifier
106                .clone()
107                .classify_response(&response_with_status(200)),
108            ClassifiedResponse::Ready(Ok(())),
109        ));
110
111        assert!(matches!(
112            classifier
113                .clone()
114                .classify_response(&response_with_status(400)),
115            ClassifiedResponse::Ready(Err(StatusInRangeFailureClass::StatusCode(
116                StatusCode::BAD_REQUEST
117            ))),
118        ));
119
120        assert!(matches!(
121            classifier.classify_response(&response_with_status(500)),
122            ClassifiedResponse::Ready(Err(StatusInRangeFailureClass::StatusCode(
123                StatusCode::INTERNAL_SERVER_ERROR
124            ))),
125        ));
126    }
127
128    fn response_with_status(status: u16) -> Response<()> {
129        Response::builder().status(status).body(()).unwrap()
130    }
131}