lago_client/
client.rs

1use reqwest::{Client as HttpClient, Response};
2use serde::de::DeserializeOwned;
3use std::time::Instant;
4use tokio::time::sleep;
5
6use lago_types::{error::LagoError, error::Result};
7
8use crate::{Config, RetryMode};
9
10/// The main client for interacting with the Lago API
11///
12/// This client handles HTTP requests, authentication, retries, and error handling
13/// when communicating with the Lago billing API.
14#[derive(Clone)]
15pub struct LagoClient {
16    pub(crate) config: Config,
17    http_client: HttpClient,
18}
19
20impl LagoClient {
21    /// Creates a new Lago client with the provided configuration
22    ///
23    /// # Arguments
24    /// * `config` - The configuration settings for the client
25    ///
26    /// # Returns
27    /// A new instance of `LagoClient`
28    pub fn new(config: Config) -> Self {
29        let http_client = HttpClient::builder()
30            .timeout(config.timeout())
31            .user_agent(config.user_agent())
32            .build()
33            .expect("Failed to create HTTP client");
34
35        Self {
36            config,
37            http_client,
38        }
39    }
40
41    /// Creates a new Lago client using default configuration from environment variables
42    ///
43    /// This method will use default settings and attempt to load credentials
44    /// from environment variables.
45    ///
46    /// # Returns
47    /// A `Result` containing a new `LagoClient` instance or an error
48    pub fn from_env() -> Result<Self> {
49        let config = Config::default();
50        Ok(Self::new(config))
51    }
52
53    /// Makes an HTTP request to the Lago API with automatic retry logic
54    ///
55    /// This method handles authentication, request serialization, response deserialization,
56    /// error handling, and automatic retries based on the configured retry policy.
57    ///
58    /// # Arguments
59    /// * `method` - The HTTP method (GET, POST, PUT, DELETE)
60    /// * `url` - The full URL to make the request to
61    /// * `body` - Optional request body that will be serialized as JSON
62    ///
63    /// # Returns
64    /// A `Result` containing the deserialized response or an error
65    pub(crate) async fn make_request<T, B>(
66        &self,
67        method: &str,
68        url: &str,
69        body: Option<&B>,
70    ) -> Result<T>
71    where
72        T: DeserializeOwned,
73        B: serde::Serialize,
74    {
75        let credentials = self.config.credentials()?;
76        let mut attempt = 0;
77
78        loop {
79            let _start_time = Instant::now();
80
81            let mut request_builder = match method {
82                "GET" => self.http_client.get(url),
83                "POST" => self.http_client.post(url),
84                "PUT" => self.http_client.put(url),
85                "DELETE" => self.http_client.delete(url),
86                _ => {
87                    return Err(LagoError::Configuration(format!(
88                        "Unsupported method: {method}"
89                    )));
90                }
91            };
92
93            request_builder = request_builder.bearer_auth(credentials.api_key());
94
95            if let Some(body) = body {
96                request_builder = request_builder.json(body);
97            }
98
99            let response = match request_builder.send().await {
100                Ok(response) => response,
101                Err(e) => {
102                    if attempt >= self.config.retry_config().max_attempts {
103                        return Err(LagoError::Http(e));
104                    }
105
106                    attempt += 1;
107                    let delay = self.config.retry_config().delay_for_attempt(attempt);
108                    sleep(delay).await;
109                    continue;
110                }
111            };
112
113            match self.handle_response(response).await {
114                Ok(result) => return Ok(result),
115                Err(e) => {
116                    if !self.should_retry(&e, attempt) {
117                        return Err(e);
118                    }
119
120                    attempt += 1;
121                    let delay = self.config.retry_config().delay_for_attempt(attempt);
122                    sleep(delay).await;
123                    continue;
124                }
125            }
126        }
127    }
128
129    /// Processes the HTTP response and converts it to the expected type
130    ///
131    /// This method handles different HTTP status codes and converts them to appropriate
132    /// error types for the client to handle.
133    ///
134    /// # Arguments
135    /// * `response` - The HTTP response from the request
136    ///
137    /// # Returns
138    /// A `Result` containing the deserialized response data or an error
139    async fn handle_response<T: DeserializeOwned>(&self, response: Response) -> Result<T> {
140        let status = response.status();
141
142        if status.is_success() {
143            let text = response.text().await.map_err(LagoError::Http)?;
144            serde_json::from_str(&text).map_err(LagoError::Serialization)
145        } else {
146            let error_text = response
147                .text()
148                .await
149                .unwrap_or_else(|_| "Unknown error".to_string());
150
151            match status.as_u16() {
152                401 => Err(LagoError::Unauthorized),
153                404 => Err(LagoError::Api {
154                    status: status.as_u16(),
155                    message: error_text,
156                }),
157                429 => Err(LagoError::RateLimit),
158                _ => Err(LagoError::Api {
159                    status: status.as_u16(),
160                    message: error_text,
161                }),
162            }
163        }
164    }
165
166    /// Determines whether a request should be retried based on the error type and attempt count
167    ///
168    /// This method implements the retry logic based on the configured retry policy.
169    /// It will retry on certain error types (HTTP errors, rate limits, server errors)
170    /// but not on client errors or authentication failures.
171    ///
172    /// # Arguments
173    /// * `error` - The error that occurred during the request
174    /// * `attempt` - The current attempt number (0-based)
175    ///
176    /// # Returns
177    /// `true` if the request should be retried, `false` otherwise
178    fn should_retry(&self, error: &LagoError, attempt: u32) -> bool {
179        if attempt >= self.config.retry_config().max_attempts {
180            return false;
181        }
182
183        if self.config.retry_config().mode == RetryMode::Off {
184            return false;
185        }
186
187        match error {
188            LagoError::Http(_) => true,
189            LagoError::RateLimit => true,
190            LagoError::Api { status, .. } => *status >= 500,
191            _ => false,
192        }
193    }
194}
195
196#[cfg(test)]
197mod tests {
198    use super::*;
199    use crate::{Config, Credentials, Region, RetryConfig, RetryMode};
200    use lago_types::error::LagoError;
201    use mockito::Server;
202    use serde::{Deserialize, Serialize};
203    use serde_json::json;
204    use std::time::Duration;
205
206    #[derive(Debug, Deserialize, Serialize)]
207    struct TestResponse {
208        id: String,
209        name: String,
210    }
211
212    #[derive(Serialize)]
213    struct TestRequest {
214        name: String,
215    }
216
217    fn create_test_client(base_url: &str) -> LagoClient {
218        let config = Config::builder()
219            .credentials(Credentials::new("test-api-key".to_string()))
220            .region(Region::Custom(base_url.to_string()))
221            .timeout(Duration::from_secs(10))
222            .build();
223
224        LagoClient::new(config)
225    }
226
227    fn create_retry_client(base_url: &str, max_attempts: u32) -> LagoClient {
228        let retry_config = RetryConfig::builder()
229            .max_attempts(max_attempts)
230            .mode(RetryMode::Adaptive)
231            .build();
232
233        let config = Config::builder()
234            .credentials(Credentials::new("test-api-key".to_string()))
235            .region(Region::Custom(base_url.to_string()))
236            .retry_config(retry_config)
237            .timeout(Duration::from_secs(5))
238            .build();
239
240        LagoClient::new(config)
241    }
242
243    #[test]
244    fn test_new_client_creation() {
245        let config = Config::default();
246        let client = LagoClient::new(config.clone());
247
248        assert_eq!(client.config.timeout(), config.timeout());
249        assert_eq!(client.config.user_agent(), config.user_agent());
250    }
251
252    #[test]
253    fn test_from_env_client_creation() {
254        let result = LagoClient::from_env();
255        assert!(result.is_ok());
256    }
257
258    #[tokio::test]
259    async fn test_successful_get_request() {
260        let mut server = Server::new_async().await;
261        let mock = server
262            .mock("GET", "/test")
263            .with_status(200)
264            .with_header("content-type", "application/json")
265            .with_body(
266                json!({
267                    "id": "123",
268                    "name": "Test"
269                })
270                .to_string(),
271            )
272            .create_async()
273            .await;
274
275        let client = create_test_client(&server.url());
276        let url = format!("{}/test", server.url());
277
278        let result: Result<TestResponse> = client.make_request("GET", &url, None::<&()>).await;
279
280        assert!(result.is_ok());
281
282        let response = result.unwrap();
283        assert_eq!(response.id, "123");
284        assert_eq!(response.name, "Test");
285
286        mock.assert_async().await;
287    }
288
289    #[tokio::test]
290    async fn test_successful_post_request() {
291        let mut server = Server::new_async().await;
292        let mock = server
293            .mock("POST", "/test")
294            .with_status(201)
295            .with_header("content-type", "application/json")
296            .match_body(mockito::Matcher::Json(json!({
297                "name": "New Item"
298            })))
299            .with_body(
300                json!({
301                    "id": "456",
302                    "name": "New Item"
303
304                })
305                .to_string(),
306            )
307            .create_async()
308            .await;
309
310        let client = create_test_client(&server.url());
311        let url = format!("{}/test", server.url());
312        let request = TestRequest {
313            name: "New Item".to_string(),
314        };
315
316        let result: Result<TestResponse> = client.make_request("POST", &url, Some(&request)).await;
317
318        assert!(result.is_ok());
319
320        let response = result.unwrap();
321        assert_eq!(response.id, "456");
322        assert_eq!(response.name, "New Item");
323
324        mock.assert_async().await;
325    }
326
327    #[tokio::test]
328    async fn test_authentication_header() {
329        let mut server = Server::new_async().await;
330        let mock = server
331            .mock("GET", "/test")
332            .match_header("Authorization", "Bearer test-api-key")
333            .with_status(200)
334            .with_body(json!({"id": "123", "name": "Test"}).to_string())
335            .create_async()
336            .await;
337
338        let client = create_test_client(&server.url());
339        let url = format!("{}/test", server.url());
340
341        let _result: Result<TestResponse> = client.make_request("GET", &url, None::<&()>).await;
342
343        mock.assert_async().await;
344    }
345
346    #[tokio::test]
347    async fn test_unsupported_method() {
348        let server = Server::new_async().await;
349        let client = create_test_client(&server.url());
350        let url = format!("{}/test", server.url());
351
352        let result: Result<TestResponse> = client.make_request("PATCH", &url, None::<&()>).await;
353
354        assert!(result.is_err());
355
356        match result.unwrap_err() {
357            LagoError::Configuration(msg) => {
358                assert!(msg.contains("Unsupported method: PATCH"));
359            }
360            _ => panic!("Expected Configuration error"),
361        }
362    }
363
364    #[tokio::test]
365    async fn test_unauthorized_error() {
366        let mut server = Server::new_async().await;
367        let mock = server
368            .mock("GET", "/test")
369            .with_status(401)
370            .with_body("Unauthorized")
371            .create_async()
372            .await;
373
374        let client = create_test_client(&server.url());
375        let url = format!("{}/test", server.url());
376
377        let result: Result<TestResponse> = client.make_request("GET", &url, None::<&()>).await;
378
379        assert!(result.is_err());
380
381        match result.unwrap_err() {
382            LagoError::Unauthorized => {}
383            _ => panic!("Expected Unauthorized error"),
384        }
385
386        mock.assert_async().await;
387    }
388
389    #[tokio::test]
390    async fn test_not_found_error() {
391        let mut server = Server::new_async().await;
392        let mock = server
393            .mock("GET", "/test")
394            .with_status(404)
395            .with_body("Not Found")
396            .create_async()
397            .await;
398
399        let client = create_test_client(&server.url());
400        let url = format!("{}/test", server.url());
401
402        let result: Result<TestResponse> = client.make_request("GET", &url, None::<&()>).await;
403
404        assert!(result.is_err());
405
406        match result.unwrap_err() {
407            LagoError::Api { status, message } => {
408                assert_eq!(status, 404);
409                assert_eq!(message, "Not Found");
410            }
411            _ => panic!("Expected Api Error"),
412        }
413
414        mock.assert_async().await;
415    }
416
417    #[tokio::test]
418    async fn test_rate_limit_error() {
419        let mut server = Server::new_async().await;
420        let mock = server
421            .mock("GET", "/test")
422            .with_status(429)
423            .with_body("Rate Limited")
424            .create_async()
425            .await;
426
427        let client = create_test_client(&server.url());
428        let url = format!("{}/test", server.url());
429
430        let result: Result<TestResponse> = client.make_request("GET", &url, None::<&()>).await;
431
432        assert!(result.is_err());
433
434        match result.unwrap_err() {
435            LagoError::RateLimit => {}
436            _ => panic!("Expected RateLimit error"),
437        }
438
439        mock.assert_async().await;
440    }
441
442    #[tokio::test]
443    async fn test_server_error() {
444        let mut server = Server::new_async().await;
445        let mock = server
446            .mock("GET", "/test")
447            .with_status(500)
448            .with_body("Internal Server Error")
449            .create_async()
450            .await;
451
452        let client = create_test_client(&server.url());
453        let url = format!("{}/test", server.url());
454
455        let result: Result<TestResponse> = client.make_request("GET", &url, None::<&()>).await;
456
457        assert!(result.is_err());
458
459        match result.unwrap_err() {
460            LagoError::Api { status, message } => {
461                assert_eq!(status, 500);
462                assert_eq!(message, "Internal Server Error");
463            }
464            _ => panic!("Expected Api Error"),
465        }
466
467        mock.assert_async().await;
468    }
469
470    #[tokio::test]
471    async fn test_json_deserialization_error() {
472        let mut server = Server::new_async().await;
473        let mock = server
474            .mock("GET", "/test")
475            .with_status(200)
476            .with_body("invalid json")
477            .create_async()
478            .await;
479
480        let client = create_test_client(&server.url());
481        let url = format!("{}/test", server.url());
482
483        let result: Result<TestResponse> = client.make_request("GET", &url, None::<&()>).await;
484
485        assert!(result.is_err());
486
487        match result.unwrap_err() {
488            LagoError::Serialization(_) => {}
489            _ => panic!("Expected Serialization error"),
490        }
491
492        mock.assert_async().await;
493    }
494
495    #[tokio::test]
496    async fn test_retry_on_server_error() {
497        let mut server = Server::new_async().await;
498        let mock = server
499            .mock("GET", "/test")
500            .with_status(500)
501            .with_body("Server Error")
502            .expect(4)
503            .create_async()
504            .await;
505
506        let client = create_retry_client(&server.url(), 3);
507        let url = format!("{}/test", server.url());
508
509        let result: Result<TestResponse> = client.make_request("GET", &url, None::<&()>).await;
510
511        assert!(result.is_err());
512        mock.assert_async().await;
513    }
514
515    #[tokio::test]
516    async fn test_retry_then_success() {
517        let mut server = Server::new_async().await;
518
519        let mock_fail = server
520            .mock("GET", "/test")
521            .with_status(500)
522            .with_body("Server Error")
523            .expect(2)
524            .create_async()
525            .await;
526
527        let mock_success = server
528            .mock("GET", "/test")
529            .with_status(200)
530            .with_body(json!({"id": "123", "name": "Success"}).to_string())
531            .expect(1)
532            .create_async()
533            .await;
534
535        let client = create_retry_client(&server.url(), 5);
536        let url = format!("{}/test", server.url());
537
538        let result: Result<TestResponse> = client.make_request("GET", &url, None::<&()>).await;
539
540        assert!(result.is_ok());
541
542        let response = result.unwrap();
543        assert_eq!(response.id, "123");
544        assert_eq!(response.name, "Success");
545
546        mock_fail.assert_async().await;
547        mock_success.assert_async().await;
548    }
549
550    #[tokio::test]
551    async fn test_no_retry_on_client_error() {
552        let mut server = Server::new_async().await;
553        let mock = server
554            .mock("GET", "/test")
555            .with_status(400)
556            .with_body("Bad Request")
557            .expect(1)
558            .create_async()
559            .await;
560
561        let client = create_retry_client(&server.url(), 3);
562        let url = format!("{}/test", server.url());
563
564        let result: Result<TestResponse> = client.make_request("GET", &url, None::<&()>).await;
565
566        assert!(result.is_err());
567        mock.assert_async().await;
568    }
569
570    #[tokio::test]
571    async fn test_should_retry_logic() {
572        let client = create_retry_client("http://localhost:8080", 3);
573
574        let rate_limit_error = LagoError::RateLimit;
575        assert!(client.should_retry(&rate_limit_error, 1));
576
577        let server_error = LagoError::Api {
578            status: 500,
579            message: "Server Error".to_string(),
580        };
581        assert!(client.should_retry(&server_error, 1));
582
583        let client_error = LagoError::Api {
584            status: 400,
585            message: "Bad Request".to_string(),
586        };
587        assert!(!client.should_retry(&client_error, 1));
588
589        assert!(!client.should_retry(&server_error, 3));
590
591        let auth_error = LagoError::Unauthorized;
592        assert!(!client.should_retry(&auth_error, 1));
593
594        let client_no_retry = create_test_client("http://localhost:8080");
595        assert!(!client_no_retry.should_retry(&server_error, 1));
596    }
597
598    #[tokio::test]
599    async fn test_timeout_handling() {
600        // Test timeout by using an unreachable address that will cause a timeout
601        let config = Config::builder()
602            .credentials(Credentials::new("test-api-key".to_string()))
603            .region(Region::Custom("http://10.255.255.1:80".to_string()))
604            .timeout(Duration::from_millis(100))
605            .build();
606
607        let client = LagoClient::new(config);
608        let url = "http://10.255.255.1:80/test";
609
610        let result: Result<TestResponse> = client.make_request("GET", url, None::<&()>).await;
611
612        assert!(result.is_err());
613
614        match result.unwrap_err() {
615            LagoError::Http(_) => {}
616            other => panic!("Expected HTTP error, got: {other:?}"),
617        }
618    }
619}