Skip to main content

cli_speedtest/
utils.rs

1// src/utils.rs
2
3use 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    // --- calculate_mbps ---
67
68    #[test]
69    fn mbps_correct_for_known_value() {
70        // 12.5 MiB in 1 second = exactly 100 Mbps
71        let bytes = 13_107_200u64; // 12.5 * 1024 * 1024
72        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    // --- with_retry ---
96
97    #[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        // max_retries = 2 means 3 total attempts: attempt 0, 1, 2
146        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}