1use crate::models::AppConfig;
4use indicatif::{ProgressBar, ProgressStyle};
5use std::time::Duration;
6use tracing::debug;
7
8pub const WARMUP_SECS: f64 = 2.0;
9
10pub fn create_spinner(msg: &str, config: &AppConfig, style_template: &str) -> ProgressBar {
11 if config.quiet {
12 ProgressBar::hidden()
13 } else {
14 let pb = ProgressBar::new_spinner();
15 if let Ok(style) = ProgressStyle::default_spinner().template(style_template) {
16 pb.set_style(style);
17 }
18 pb.set_message(msg.to_string());
19 pb.enable_steady_tick(Duration::from_millis(100));
20 pb
21 }
22}
23
24pub fn calculate_mbps(bytes: u64, duration_secs: f64) -> f64 {
25 if duration_secs <= 0.0 {
26 return 0.0;
27 }
28 let megabytes = (bytes as f64) / (1024.0 * 1024.0);
29 (megabytes * 8.0) / duration_secs
30}
31
32pub async fn with_retry<F, Fut, T>(max_retries: u32, mut f: F) -> anyhow::Result<T>
33where
34 F: FnMut() -> Fut,
35 Fut: std::future::Future<Output = anyhow::Result<T>>,
36{
37 let mut last_err = anyhow::anyhow!("No attempts made");
38 for attempt in 0..=max_retries {
39 match f().await {
40 Ok(val) => return Ok(val),
41 Err(e) => {
42 if attempt < max_retries {
43 let backoff = Duration::from_millis(100 * 2u64.pow(attempt));
44 debug!(
45 "Request failed (attempt {}/{}): {}. Retrying in {:?}...",
46 attempt + 1,
47 max_retries + 1,
48 e,
49 backoff
50 );
51 tokio::time::sleep(backoff).await;
52 }
53 last_err = e;
54 }
55 }
56 }
57 Err(last_err)
58}
59
60#[cfg(test)]
61mod tests {
62 use super::*;
63 use std::sync::Arc;
64 use std::sync::atomic::{AtomicU32, Ordering};
65
66 #[test]
69 fn mbps_correct_for_known_value() {
70 let bytes = 13_107_200u64; let speed = calculate_mbps(bytes, 1.0);
73 assert!(
74 (speed - 100.0).abs() < 0.001,
75 "Expected 100 Mbps, got {}",
76 speed
77 );
78 }
79
80 #[test]
81 fn mbps_zero_for_zero_duration() {
82 assert_eq!(calculate_mbps(1_000_000, 0.0), 0.0);
83 }
84
85 #[test]
86 fn mbps_zero_for_negative_duration() {
87 assert_eq!(calculate_mbps(1_000_000, -5.0), 0.0);
88 }
89
90 #[test]
91 fn mbps_zero_bytes_gives_zero() {
92 assert_eq!(calculate_mbps(0, 10.0), 0.0);
93 }
94
95 #[tokio::test]
98 async fn retry_succeeds_on_first_attempt() {
99 let result = with_retry(3, || async { Ok::<i32, anyhow::Error>(42) }).await;
100 assert!(result.is_ok());
101 assert_eq!(result.unwrap(), 42);
102 }
103
104 #[tokio::test]
105 async fn retry_succeeds_on_second_attempt() {
106 let attempts = Arc::new(AtomicU32::new(0));
107 let attempts_c = attempts.clone();
108
109 let result = with_retry(3, move || {
110 let counter = attempts_c.clone();
111 async move {
112 let n = counter.fetch_add(1, Ordering::SeqCst);
113 if n == 0 {
114 anyhow::bail!("transient error");
115 }
116 Ok::<i32, anyhow::Error>(99)
117 }
118 })
119 .await;
120
121 assert!(result.is_ok());
122 assert_eq!(result.unwrap(), 99);
123 assert_eq!(
124 attempts.load(Ordering::SeqCst),
125 2,
126 "Should have taken exactly 2 attempts"
127 );
128 }
129
130 #[tokio::test]
131 async fn retry_exhausts_all_attempts_and_returns_last_error() {
132 let attempts = Arc::new(AtomicU32::new(0));
133 let attempts_c = attempts.clone();
134
135 let result: anyhow::Result<()> = with_retry(2, move || {
136 let counter = attempts_c.clone();
137 async move {
138 counter.fetch_add(1, Ordering::SeqCst);
139 anyhow::bail!("always fails")
140 }
141 })
142 .await;
143
144 assert!(result.is_err());
145 assert_eq!(
147 attempts.load(Ordering::SeqCst),
148 3,
149 "Should have attempted exactly max_retries + 1 times"
150 );
151 }
152
153 #[tokio::test]
154 async fn retry_with_zero_retries_attempts_exactly_once() {
155 let attempts = Arc::new(AtomicU32::new(0));
156 let attempts_c = attempts.clone();
157
158 let result: anyhow::Result<()> = with_retry(0, move || {
159 let counter = attempts_c.clone();
160 async move {
161 counter.fetch_add(1, Ordering::SeqCst);
162 anyhow::bail!("fail")
163 }
164 })
165 .await;
166
167 assert!(result.is_err());
168 assert_eq!(
169 attempts.load(Ordering::SeqCst),
170 1,
171 "Zero retries = exactly 1 attempt"
172 );
173 }
174}