brainwires_providers/
http_client.rs1use reqwest::{Client, RequestBuilder};
7use std::sync::Arc;
8
9use super::rate_limiter::RateLimiter;
10
11#[derive(Clone, Debug)]
13pub struct RateLimitedClient {
14 client: Client,
15 limiter: Option<Arc<RateLimiter>>,
16}
17
18impl RateLimitedClient {
19 pub fn new() -> Self {
21 Self {
22 client: Client::new(),
23 limiter: None,
24 }
25 }
26
27 pub fn with_rate_limit(requests_per_minute: u32) -> Self {
29 Self {
30 client: Client::new(),
31 limiter: Some(Arc::new(RateLimiter::new(requests_per_minute))),
32 }
33 }
34
35 pub fn from_client(client: Client, requests_per_minute: Option<u32>) -> Self {
37 Self {
38 client,
39 limiter: requests_per_minute.map(|rpm| Arc::new(RateLimiter::new(rpm))),
40 }
41 }
42
43 pub fn inner(&self) -> &Client {
45 &self.client
46 }
47
48 pub async fn get(&self, url: &str) -> RequestBuilder {
50 self.wait_for_token().await;
51 self.client.get(url)
52 }
53
54 pub async fn post(&self, url: &str) -> RequestBuilder {
56 self.wait_for_token().await;
57 self.client.post(url)
58 }
59
60 async fn wait_for_token(&self) {
62 if let Some(ref limiter) = self.limiter {
63 limiter.acquire().await;
64 }
65 }
66
67 pub fn available_tokens(&self) -> Option<u32> {
70 self.limiter.as_ref().map(|l| l.available_tokens())
71 }
72}
73
74impl Default for RateLimitedClient {
75 fn default() -> Self {
76 Self::new()
77 }
78}
79
80#[cfg(test)]
81mod tests {
82 use super::*;
83
84 #[test]
85 fn test_default_client() {
86 let client = RateLimitedClient::new();
87 assert!(client.available_tokens().is_none());
88 }
89
90 #[test]
91 fn test_rate_limited_client() {
92 let client = RateLimitedClient::with_rate_limit(60);
93 assert_eq!(client.available_tokens(), Some(60));
94 }
95
96 #[tokio::test]
97 async fn test_post_consumes_token() {
98 let client = RateLimitedClient::with_rate_limit(10);
99 let _req = client.post("https://example.com").await;
100 assert_eq!(client.available_tokens(), Some(9));
101 }
102}