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#[derive(Clone)]
15pub struct LagoClient {
16 pub(crate) config: Config,
17 http_client: HttpClient,
18}
19
20impl LagoClient {
21 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 pub fn from_env() -> Result<Self> {
49 let config = Config::default();
50 Ok(Self::new(config))
51 }
52
53 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 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 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 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}