ddclient_rs/
rate.rs

1// Copyright (c) 2023, Direct Decisions Rust client AUTHORS.
2// All rights reserved.
3// Use of this source code is governed by a BSD-style
4// license that can be found in the LICENSE file.
5
6const HEADER_RATE_LIMIT: &str = "X-RateLimit-Limit";
7const HEADER_RATE_REMAINING: &str = "X-RateLimit-Remaining";
8const HEADER_RATE_RESET: &str = "X-RateLimit-Reset";
9const HEADER_RATE_RETRY: &str = "Retry-After";
10
11use reqwest::header::HeaderMap;
12use std::str::FromStr;
13use std::time::{Duration, SystemTime, UNIX_EPOCH};
14
15/// Represents the rate limit information returned by the API.
16///
17/// This struct contains the rate limit information returned by the API, including the number of
18/// requests allowed, the number of requests remaining, and the time at which the rate limit will
19/// reset.
20///
21#[derive(Clone, Default, Debug)]
22pub struct Rate {
23    pub limit: u32,
24    pub remaining: u32,
25    pub reset: u64,
26    pub retry: u64,
27}
28
29impl Rate {
30    pub(crate) fn from_headers(headers: &HeaderMap) -> Option<Self> {
31        let limit = fetch_header(headers, HEADER_RATE_LIMIT)?;
32        let remaining = fetch_header(headers, HEADER_RATE_REMAINING)?;
33        let reset_secs: u64 = fetch_header(headers, HEADER_RATE_RESET)?;
34        let retry_secs: u64 = fetch_header(headers, HEADER_RATE_RETRY)?;
35
36        let now = SystemTime::now().duration_since(UNIX_EPOCH).ok()?;
37        let reset = now + Duration::from_secs(reset_secs);
38        let retry = now + Duration::from_secs(retry_secs);
39
40        Some(Self {
41            limit,
42            remaining,
43            reset: reset.as_secs(),
44            retry: retry.as_secs(),
45        })
46    }
47}
48
49fn fetch_header<T>(headers: &HeaderMap, header: &str) -> Option<T>
50where
51    T: FromStr,
52{
53    headers.get(header)?.to_str().ok()?.parse::<T>().ok()
54}
55
56// rate tests
57#[cfg(test)]
58mod tests {
59    use super::*;
60    use httpmock::Method::GET;
61    use httpmock::MockServer;
62
63    #[tokio::test]
64    async fn test_rate_from_headers() {
65        let server = MockServer::start();
66        let mock = server.mock(|when, then| {
67            when.method(GET).path("/test");
68            then.status(200)
69                .header(HEADER_RATE_LIMIT, "100")
70                .header(HEADER_RATE_REMAINING, "50")
71                .header(HEADER_RATE_RESET, "1000")
72                .header(HEADER_RATE_RETRY, "1000");
73        });
74
75        let client = reqwest::Client::new();
76        let response = client.get(server.url("/test")).send().await.unwrap();
77        let rate = Rate::from_headers(response.headers()).unwrap();
78
79        assert_eq!(rate.limit, 100);
80        assert_eq!(rate.remaining, 50);
81        // assert that rate is time.now + reset from seconds rounded to 5 seconds
82
83        let now = SystemTime::now().duration_since(UNIX_EPOCH).ok().unwrap();
84        let reset = now + Duration::from_secs(1000);
85        assert_eq!(rate.reset, reset.as_secs());
86        let retry = now + Duration::from_secs(1000);
87        assert_eq!(rate.retry, retry.as_secs());
88        mock.assert();
89    }
90
91    #[tokio::test]
92    async fn test_no_headers() {
93        let server = MockServer::start();
94        let mock = server.mock(|when, then| {
95            when.method(GET).path("/test");
96            then.status(200);
97        });
98
99        let client = reqwest::Client::new();
100        let response = client.get(server.url("/test")).send().await.unwrap();
101        let rate = Rate::from_headers(response.headers());
102        assert!(rate.is_none());
103        mock.assert();
104    }
105}