rate_limits/retryafter/
mod.rs1use std::str::FromStr;
5
6use headers::HeaderValue;
7use time::{Date, format_description::well_known::Rfc2822};
8
9use crate::{
10 casesensitive_headermap::CaseSensitiveHeaderMap,
11 reset_time::{ResetTime, ResetTimeKind},
12};
13
14use super::error::{Error, Result};
15
16#[derive(Copy, Clone, Debug, PartialEq)]
18pub struct RateLimit {
19 pub reset: ResetTime,
21}
22
23impl RateLimit {
24 pub fn new<T: Into<CaseSensitiveHeaderMap>>(headers: T) -> std::result::Result<Self, Error> {
28 let headers = headers.into();
29 let reset = match Self::get_retry_after_header(&headers) {
30 Some(retry_after) => {
31 if Date::parse(retry_after.to_str()?, &Rfc2822).is_ok() {
32 ResetTime::new(retry_after, ResetTimeKind::ImfFixdate)?
33 } else {
34 ResetTime::new(retry_after, ResetTimeKind::Seconds)?
35 }
36 }
37 None => return Err(Error::MissingRetryAfter),
38 };
39
40 Ok(RateLimit { reset })
41 }
42
43 fn get_retry_after_header(header_map: &CaseSensitiveHeaderMap) -> Option<&HeaderValue> {
48 header_map
49 .get("Retry-After")
50 .or_else(|| header_map.get("retry-after"))
51 }
52
53 #[must_use]
55 pub const fn reset(&self) -> ResetTime {
56 self.reset
57 }
58}
59
60impl FromStr for RateLimit {
61 type Err = Error;
62
63 fn from_str(map: &str) -> Result<Self> {
64 RateLimit::new(CaseSensitiveHeaderMap::from_str(map)?)
65 }
66}
67
68#[cfg(test)]
69mod tests {
70
71 use super::*;
72 use indoc::indoc;
73 use time::macros::datetime;
74
75 #[test]
76 fn parse_retry_after_seconds() {
77 let map = CaseSensitiveHeaderMap::from_str("Retry-After: 30").unwrap();
78 let retry = RateLimit::get_retry_after_header(&map).unwrap();
79
80 assert_eq!("30", retry);
81 }
82
83 #[test]
84 fn retry_after_seconds() {
85 let headers = indoc! {"
86 Retry-After: 19
87 "};
88
89 let rate = RateLimit::from_str(headers).unwrap();
90 assert_eq!(rate.reset(), ResetTime::Seconds(19));
91 }
92
93 #[test]
94 fn retry_after_seconds_case_sensitive() {
95 let headers = indoc! {"
96 retry-after: 19
97 "};
98
99 let rate = RateLimit::from_str(headers).unwrap();
100 assert_eq!(rate.reset(), ResetTime::Seconds(19));
101 }
102
103 #[test]
104 fn retry_after_imf_fixdate() {
105 let headers = indoc! {"
106 Retry-After: Fri, 31 Dec 1999 23:59:59 GMT
107 "};
108
109 let rate = RateLimit::from_str(headers).unwrap();
110 assert_eq!(
111 rate.reset(),
112 ResetTime::DateTime(datetime!(1999-12-31 23:59:59 UTC))
113 );
114 }
115}