use super::{Action, Attempt, Policy};
#[derive(Clone, Copy, Debug)]
pub struct Limited {
remaining: usize,
}
impl Limited {
pub fn new(max: usize) -> Self {
Limited { remaining: max }
}
}
impl Default for Limited {
fn default() -> Self {
Limited::new(20)
}
}
impl<B, E> Policy<B, E> for Limited {
fn redirect(&mut self, _: &Attempt<'_>) -> Result<Action, E> {
if self.remaining > 0 {
self.remaining -= 1;
Ok(Action::Follow)
} else {
Ok(Action::Stop)
}
}
}
#[cfg(test)]
mod tests {
use http::{Request, Uri};
use super::*;
#[test]
fn works() {
let uri = Uri::from_static("https://example.com/");
let mut policy = Limited::new(2);
for _ in 0..2 {
let mut request = Request::builder().uri(uri.clone()).body(()).unwrap();
Policy::<(), ()>::on_request(&mut policy, &mut request);
let attempt = Attempt {
status: Default::default(),
location: &uri,
previous: &uri,
};
assert!(Policy::<(), ()>::redirect(&mut policy, &attempt)
.unwrap()
.is_follow());
}
let mut request = Request::builder().uri(uri.clone()).body(()).unwrap();
Policy::<(), ()>::on_request(&mut policy, &mut request);
let attempt = Attempt {
status: Default::default(),
location: &uri,
previous: &uri,
};
assert!(Policy::<(), ()>::redirect(&mut policy, &attempt)
.unwrap()
.is_stop());
}
}