azoth_balancer/
config.rs

1use serde::Deserialize;
2use std::{collections::HashSet, fs};
3use thiserror::Error;
4use tracing::{error, info, warn};
5
6fn resolve_env_vars_in_urls(endpoints: &mut [EndpointConfig]) {
7    use std::env;
8
9    for endpoint in endpoints {
10        if endpoint.url.starts_with("${") && endpoint.url.ends_with('}') {
11            let var_name = endpoint.url[2..endpoint.url.len() - 1].to_string();
12
13            if let Ok(value) = env::var(&var_name) {
14                endpoint.url = value;
15                info!(var_name = var_name, "Resolved endpoint URL from environment variable");
16            } else {
17                warn!(var_name = var_name, "Environment variable not found for endpoint URL");
18            }
19        }
20    }
21}
22
23#[derive(Debug, Error)]
24pub enum ConfigError {
25    #[error("Configuration error: {0}")]
26    ConfigError(String),
27}
28
29#[derive(Debug, Deserialize, Clone, Default)]
30pub struct Config {
31    pub server: Option<ServerConfig>,
32    pub balancer: Option<BalancerConfig>,
33}
34
35impl Config {
36    /// Applies defaults, validates, and sanitizes the configuration.
37    /// This ensures that the configuration is in a consistent and usable state
38    /// by filling in missing values and clamping others to valid ranges.
39    pub fn finalize(mut self) -> Result<Self, ConfigError> {
40        let mut server_cfg = self.server.take().unwrap_or_default();
41        server_cfg.bind_addr = server_cfg.bind_addr.or_else(|| Some(DEFAULT_BIND_ADDR.to_string()));
42        self.server = Some(server_cfg);
43
44        let mut balancer_cfg = self.balancer.take().unwrap_or_default();
45
46        balancer_cfg.health_check_interval_secs =
47            balancer_cfg.health_check_interval_secs.or(Some(DEFAULT_HEALTH_CHECK_INTERVAL_SECS));
48        balancer_cfg.health_check_timeout_secs =
49            balancer_cfg.health_check_timeout_secs.or(Some(DEFAULT_HEALTH_CHECK_TIMEOUT_SECS));
50        balancer_cfg.base_cooldown_secs =
51            balancer_cfg.base_cooldown_secs.or(Some(DEFAULT_BASE_COOLDOWN_SECS));
52        balancer_cfg.max_cooldown_secs =
53            balancer_cfg.max_cooldown_secs.or(Some(DEFAULT_MAX_COOLDOWN_SECS));
54
55        balancer_cfg.latency_smoothing_factor = Some(
56            balancer_cfg
57                .latency_smoothing_factor
58                .unwrap_or(DEFAULT_LATENCY_SMOOTHING_FACTOR)
59                .clamp(0.0, 1.0),
60        );
61
62        balancer_cfg.max_batch_size =
63            Some(balancer_cfg.max_batch_size.unwrap_or(DEFAULT_MAX_BATCH_SIZE).max(1));
64        balancer_cfg.max_concurrency =
65            Some(balancer_cfg.max_concurrency.unwrap_or(DEFAULT_MAX_CONCURRENCY).max(1));
66
67        // Client settings
68        balancer_cfg.connect_timeout_ms =
69            balancer_cfg.connect_timeout_ms.or(Some(DEFAULT_CONNECT_TIMEOUT_MS));
70        balancer_cfg.timeout_secs = balancer_cfg.timeout_secs.or(Some(DEFAULT_TIMEOUT_SECS));
71        balancer_cfg.pool_idle_timeout_secs =
72            balancer_cfg.pool_idle_timeout_secs.or(Some(DEFAULT_POOL_IDLE_TIMEOUT_SECS));
73        balancer_cfg.pool_max_idle_per_host =
74            balancer_cfg.pool_max_idle_per_host.or(Some(DEFAULT_POOL_MAX_IDLE_PER_HOST));
75
76        let mut endpoints = balancer_cfg.endpoints.take().unwrap_or_else(get_default_endpoints);
77
78        resolve_env_vars_in_urls(&mut endpoints);
79
80        endpoints = validate_and_dedupe_endpoints(endpoints)?;
81
82        // Apply final constraints to endpoints after validation
83        for ep in &mut endpoints {
84            ep.rate_limit_per_sec = ep.rate_limit_per_sec.max(1);
85            ep.burst_size = ep.burst_size.max(1);
86            if ep.weight.is_none() {
87                ep.weight = Some(DEFAULT_ENDPOINT_WEIGHT);
88            }
89        }
90
91        balancer_cfg.endpoints = Some(endpoints);
92        self.balancer = Some(balancer_cfg);
93
94        Ok(self)
95    }
96}
97
98#[derive(Debug, Deserialize, Clone, Default)]
99pub struct ServerConfig {
100    pub bind_addr: Option<String>,
101}
102
103#[derive(Debug, Deserialize, Clone, Default)]
104pub struct BalancerConfig {
105    pub health_check_interval_secs: Option<u64>,
106    pub health_check_timeout_secs: Option<u64>,
107    pub base_cooldown_secs: Option<u64>,
108    pub max_cooldown_secs: Option<u64>,
109    pub latency_smoothing_factor: Option<f64>,
110    pub endpoints: Option<Vec<EndpointConfig>>,
111    pub max_batch_size: Option<usize>,
112    pub max_concurrency: Option<usize>,
113    // Client timeout and pool settings
114    pub connect_timeout_ms: Option<u64>,
115    pub timeout_secs: Option<u64>,
116    pub pool_idle_timeout_secs: Option<u64>,
117    pub pool_max_idle_per_host: Option<usize>,
118}
119
120#[derive(Debug, Deserialize, Clone)]
121pub struct EndpointConfig {
122    pub name: Option<String>, // Add this line - optional name
123    pub url: String,
124    pub rate_limit_per_sec: u32,
125    #[serde(default = "default_burst_size")]
126    pub burst_size: u32,
127    #[serde(default)]
128    pub weight: Option<u32>,
129}
130
131impl Default for EndpointConfig {
132    fn default() -> Self {
133        Self {
134            name: None, // Add this line
135            url: String::new(),
136            rate_limit_per_sec: DEFAULT_ENDPOINT_RATE_LIMIT,
137            burst_size: DEFAULT_BURST_SIZE,
138            weight: None,
139        }
140    }
141}
142
143// Updated defaults based on the user's provided configuration.
144pub const DEFAULT_BURST_SIZE: u32 = 10;
145pub const DEFAULT_BIND_ADDR: &str = "127.0.0.1:8549";
146pub const DEFAULT_HEALTH_CHECK_INTERVAL_SECS: u64 = 30;
147pub const DEFAULT_HEALTH_CHECK_TIMEOUT_SECS: u64 = 5;
148pub const DEFAULT_BASE_COOLDOWN_SECS: u64 = 3;
149pub const DEFAULT_MAX_COOLDOWN_SECS: u64 = 60;
150pub const DEFAULT_LATENCY_SMOOTHING_FACTOR: f64 = 0.1;
151pub const DEFAULT_ENDPOINT_RATE_LIMIT: u32 = 5;
152pub const DEFAULT_ENDPOINT_WEIGHT: u32 = 20;
153pub const DEFAULT_MAX_BATCH_SIZE: usize = 100;
154
155// Default client setting constants
156pub const DEFAULT_CONNECT_TIMEOUT_MS: u64 = 500;
157pub const DEFAULT_TIMEOUT_SECS: u64 = 5;
158pub const DEFAULT_POOL_IDLE_TIMEOUT_SECS: u64 = 25;
159pub const DEFAULT_POOL_MAX_IDLE_PER_HOST: usize = 100;
160pub const DEFAULT_MAX_CONCURRENCY: usize = 100;
161
162// A curated list of default endpoints based on user's high-quality providers.
163pub const DEFAULT_ENDPOINTS: [&str; 4] = [
164    "https://arbitrum-one-public.nodies.app",
165    "https://arbitrum-one.public.blastapi.io",
166    "https://arbitrum.meowrpc.com",
167    "https://arbitrum.drpc.org/",
168];
169
170pub fn default_burst_size() -> u32 {
171    DEFAULT_BURST_SIZE
172}
173
174pub fn try_load_config(path: &str) -> Result<Option<Config>, ConfigError> {
175    match fs::read_to_string(path) {
176        Ok(raw) => match toml::from_str::<Config>(&raw) {
177            Ok(cfg) => {
178                info!(path = %path, "Loaded config");
179                Ok(Some(cfg))
180            }
181            Err(e) => {
182                error!(path = %path, error = %e, "Failed to parse config");
183                Err(ConfigError::ConfigError(e.to_string()))
184            }
185        },
186        Err(e) => {
187            if e.kind() == std::io::ErrorKind::NotFound {
188                info!(path = %path, "No config file found, using defaults");
189                Ok(None)
190            } else {
191                Err(ConfigError::ConfigError(e.to_string()))
192            }
193        }
194    }
195}
196
197pub fn validate_and_dedupe_endpoints(
198    endpoints: Vec<EndpointConfig>,
199) -> Result<Vec<EndpointConfig>, ConfigError> {
200    let mut seen = HashSet::new();
201    const MAX_URL_LEN: usize = 2048;
202    const MAX_RATE_LIMIT: u32 = 100_000;
203    const MAX_BURST_SIZE: u32 = 10_000;
204
205    let validated_endpoints: Vec<EndpointConfig> = endpoints
206        .into_iter()
207        .filter_map(|mut e| {
208            e.url = e.url.trim().to_string();
209
210            if e.url.is_empty() {
211                warn!("Skipping empty endpoint URL");
212                return None;
213            }
214
215            if !e.url.to_lowercase().starts_with("http://") && !e.url.to_lowercase().starts_with("https://") {
216                warn!(url = %e.url, "Skipping invalid endpoint URL");
217                return None;
218            }
219
220            if e.url.len() > MAX_URL_LEN {
221                warn!(url = %e.url, "Skipping endpoint exceeding max length");
222                return None;
223            }
224
225            if e.url.chars().any(|c| c.is_control() || c.is_whitespace()) {
226                warn!(url = %e.url, "Skipping endpoint with invalid characters");
227                return None;
228            }
229
230            if e.rate_limit_per_sec == 0 || e.rate_limit_per_sec > MAX_RATE_LIMIT {
231                warn!(url = %e.url, rate = e.rate_limit_per_sec, "Skipping endpoint with invalid rate_limit_per_sec");
232                return None;
233            }
234
235            if e.burst_size == 0 || e.burst_size > MAX_BURST_SIZE {
236                warn!(url = %e.url, burst = e.burst_size, "Skipping endpoint with invalid burst_size");
237                return None;
238            }
239
240            // Canonicalize: lowercase scheme + remove trailing slash
241            let mut canonical_url = e.url.clone();
242            if canonical_url.to_lowercase().starts_with("http://") {
243                canonical_url.replace_range(0..7, "http://");
244            } else if canonical_url.to_lowercase().starts_with("https://") {
245                canonical_url.replace_range(0..8, "https://");
246            }
247            while canonical_url.ends_with('/') {
248                canonical_url.pop();
249            }
250            e.url = canonical_url;
251
252            if seen.insert(e.url.clone()) {
253                Some(e)
254            } else {
255                None
256            }
257        })
258        .collect();
259
260    if validated_endpoints.is_empty() {
261        return Err(ConfigError::ConfigError("No valid endpoints configured".to_string()));
262    }
263
264    Ok(validated_endpoints)
265}
266
267pub fn get_default_endpoints() -> Vec<EndpointConfig> {
268    DEFAULT_ENDPOINTS
269        .iter()
270        .map(|&s| EndpointConfig {
271            name: None,
272            url: s.to_string(),
273            rate_limit_per_sec: DEFAULT_ENDPOINT_RATE_LIMIT,
274            burst_size: default_burst_size(),
275            weight: Some(DEFAULT_ENDPOINT_WEIGHT),
276        })
277        .collect()
278}
279
280// Note: This is just the structure, not the full code.
281#[cfg(test)]
282mod tests {
283    use super::*;
284    use std::io::Write;
285    use tempfile::NamedTempFile;
286
287    #[test]
288    fn test_try_load_config_valid_file() {
289        let mut file = NamedTempFile::new().unwrap();
290        writeln!(file, "[server]\nbind_addr = \"127.0.0.1:8070\"").unwrap();
291        let path = file.path().to_str().unwrap();
292        let result = try_load_config(path).unwrap();
293        assert!(result.is_some());
294        let config = result.unwrap();
295        assert_eq!(config.server.unwrap().bind_addr.unwrap(), "127.0.0.1:8070");
296    }
297
298    #[test]
299    fn test_try_load_config_file_not_found() {
300        let result = try_load_config("nonexistent.toml").unwrap();
301        assert!(result.is_none());
302    }
303
304    #[test]
305    fn test_try_load_config_invalid_file() {
306        let mut file = NamedTempFile::new().unwrap();
307        writeln!(file, "[server]\nbind_addr = 12345").unwrap();
308        let path = file.path().to_str().unwrap();
309        let result = try_load_config(path);
310        assert!(result.is_err());
311    }
312
313    #[test]
314    fn test_validate_and_dedupe_endpoints() {
315        let endpoints = vec![
316            EndpointConfig {
317                name: None,
318                url: "https://valid1.com".into(),
319                rate_limit_per_sec: 10,
320                burst_size: 5,
321                weight: Some(1),
322            },
323            EndpointConfig {
324                name: None,
325                url: "https://valid1.com".into(),
326                rate_limit_per_sec: 10,
327                burst_size: 5,
328                weight: Some(1),
329            },
330            EndpointConfig {
331                name: None,
332                url: "".into(),
333                rate_limit_per_sec: 10,
334                burst_size: 5,
335                weight: Some(1),
336            },
337            EndpointConfig {
338                name: None,
339                url: "invalid_url".into(),
340                rate_limit_per_sec: 10,
341                burst_size: 5,
342                weight: Some(1),
343            },
344        ];
345        let result = validate_and_dedupe_endpoints(endpoints);
346        assert!(result.is_ok());
347        let list = result.unwrap();
348        assert_eq!(list.len(), 1);
349        assert_eq!(list[0].url, "https://valid1.com");
350    }
351
352    #[test]
353    fn test_get_default_endpoints() {
354        let endpoints = get_default_endpoints();
355        assert_eq!(endpoints.len(), 4);
356        assert!(endpoints.iter().all(|e| e.url.starts_with("https://")));
357    }
358
359    #[test]
360    fn test_validate_and_dedupe_all_invalid() {
361        let endpoints = vec![
362            EndpointConfig {
363                name: None,
364                url: "not-a-url".into(),
365                rate_limit_per_sec: 10,
366                burst_size: 5,
367                weight: Some(1),
368            },
369            EndpointConfig {
370                name: None,
371                url: "".into(),
372                rate_limit_per_sec: 10,
373                burst_size: 5,
374                weight: Some(1),
375            },
376        ];
377        let result = validate_and_dedupe_endpoints(endpoints);
378        assert!(result.is_err());
379    }
380
381    #[test]
382    fn test_validate_and_dedupe_multiple_valid() {
383        let endpoints = vec![
384            EndpointConfig {
385                name: None,
386                url: "https://mainnet.infura.io/v3/123".into(),
387                rate_limit_per_sec: 20,
388                burst_size: 5,
389                weight: Some(1),
390            },
391            EndpointConfig {
392                name: None,
393                url: "https://rpc.ankr.com/eth".into(),
394                rate_limit_per_sec: 20,
395                burst_size: 5,
396                weight: Some(1),
397            },
398        ];
399        let result = validate_and_dedupe_endpoints(endpoints).unwrap();
400        assert_eq!(result.len(), 2);
401        assert!(result.iter().any(|e| e.url.contains("infura")));
402        assert!(result.iter().any(|e| e.url.contains("ankr")));
403    }
404
405    #[test]
406    fn test_balancer_config_optional_fields() {
407        let toml = r#"
408        [balancer]
409        max_batch_size = 42
410        max_concurrency = 10
411        "#;
412
413        let mut file = NamedTempFile::new().unwrap();
414        writeln!(file, "{}", toml).unwrap();
415        let path = file.path().to_str().unwrap();
416
417        let cfg = try_load_config(path).unwrap().unwrap();
418        let balancer = cfg.balancer.unwrap();
419        assert_eq!(balancer.max_batch_size.unwrap(), 42);
420        assert_eq!(balancer.max_concurrency.unwrap(), 10);
421        assert!(balancer.base_cooldown_secs.is_none());
422        assert!(balancer.endpoints.is_none());
423    }
424
425    #[test]
426    fn test_endpoint_default_values() {
427        let endpoint = EndpointConfig {
428            url: "https://example.com".to_string(),
429            rate_limit_per_sec: 10,
430            ..Default::default()
431        };
432        assert_eq!(endpoint.burst_size, DEFAULT_BURST_SIZE);
433        assert!(endpoint.weight.is_none());
434    }
435
436    #[test]
437    fn test_get_default_endpoints_not_empty() {
438        let endpoints = get_default_endpoints();
439        assert!(!endpoints.is_empty());
440        for e in &endpoints {
441            assert!(e.url.starts_with("https://"));
442            assert_eq!(e.rate_limit_per_sec, DEFAULT_ENDPOINT_RATE_LIMIT);
443            assert_eq!(e.burst_size, DEFAULT_BURST_SIZE);
444            assert_eq!(e.weight.unwrap(), DEFAULT_ENDPOINT_WEIGHT);
445        }
446    }
447
448    #[test]
449    fn test_server_config_defaults() {
450        let cfg = Config::default();
451        assert!(cfg.server.is_none());
452    }
453
454    #[test]
455    fn test_endpoint_custom_values() {
456        let endpoint = EndpointConfig {
457            name: None,
458            url: "https://example.com".to_string(),
459            rate_limit_per_sec: 0,
460            burst_size: 50,
461            weight: Some(5),
462        };
463        assert_eq!(endpoint.rate_limit_per_sec, 0);
464        assert_eq!(endpoint.burst_size, 50);
465        assert_eq!(endpoint.weight, Some(5));
466    }
467
468    #[test]
469    fn test_validate_endpoints_trim_and_canonicalize() {
470        let endpoints = vec![
471            EndpointConfig {
472                name: None,
473                url: " HTTPS://example.com/ ".into(),
474                rate_limit_per_sec: 10,
475                burst_size: 5,
476                weight: None,
477            },
478            EndpointConfig {
479                name: None,
480                url: "https://example.com".into(),
481                rate_limit_per_sec: 10,
482                burst_size: 5,
483                weight: None,
484            },
485        ];
486        let validated = validate_and_dedupe_endpoints(endpoints).unwrap();
487        assert_eq!(validated.len(), 1);
488        assert_eq!(validated[0].url, "https://example.com");
489    }
490
491    #[test]
492    fn test_validate_and_dedupe_mixed() {
493        let endpoints = vec![
494            EndpointConfig {
495                name: None,
496                url: "https://good.com".into(),
497                rate_limit_per_sec: 10,
498                burst_size: 5,
499                weight: None,
500            },
501            EndpointConfig {
502                name: None,
503                url: "bad".into(),
504                rate_limit_per_sec: 10,
505                burst_size: 5,
506                weight: None,
507            },
508            EndpointConfig {
509                name: None,
510                url: "".into(),
511                rate_limit_per_sec: 10,
512                burst_size: 5,
513                weight: None,
514            },
515        ];
516        let validated = validate_and_dedupe_endpoints(endpoints).unwrap();
517        assert_eq!(validated.len(), 1);
518        assert_eq!(validated[0].url, "https://good.com");
519    }
520
521    #[test]
522    fn test_get_default_endpoints_exact() {
523        let endpoints = get_default_endpoints();
524        for (i, e) in endpoints.iter().enumerate() {
525            assert_eq!(e.url, DEFAULT_ENDPOINTS[i]);
526            assert_eq!(e.burst_size, DEFAULT_BURST_SIZE);
527            assert_eq!(e.rate_limit_per_sec, DEFAULT_ENDPOINT_RATE_LIMIT);
528            assert_eq!(e.weight.unwrap(), DEFAULT_ENDPOINT_WEIGHT);
529        }
530    }
531
532    #[test]
533    fn test_try_load_empty_file() {
534        let mut file = NamedTempFile::new().unwrap();
535        writeln!(file).unwrap();
536        let path = file.path().to_str().unwrap();
537        let result = try_load_config(path).unwrap();
538        assert!(result.is_some());
539        let cfg = result.unwrap();
540        assert!(cfg.server.is_none());
541        assert!(cfg.balancer.is_none());
542    }
543
544    #[test]
545    fn test_validate_endpoints_invalid_characters() {
546        let endpoints = vec![
547            EndpointConfig {
548                name: None,
549                url: "https://good.com/ bad".into(),
550                rate_limit_per_sec: 10,
551                burst_size: 5,
552                weight: None,
553            },
554            EndpointConfig {
555                name: None,
556                url: "https://ok.com/\x07".into(),
557                rate_limit_per_sec: 10,
558                burst_size: 5,
559                weight: None,
560            },
561        ];
562        let result = validate_and_dedupe_endpoints(endpoints);
563        assert!(result.is_err());
564    }
565
566    #[test]
567    fn test_validate_endpoints_max_url_length() {
568        let long_url = format!("https://example.com/{}", "a".repeat(2048));
569        let endpoints = vec![EndpointConfig {
570            name: None,
571            url: long_url.clone(),
572            rate_limit_per_sec: 10,
573            burst_size: 5,
574            weight: None,
575        }];
576        let result = validate_and_dedupe_endpoints(endpoints);
577        assert!(result.is_err(), "URL exceeding max length should be rejected");
578    }
579
580    #[test]
581    fn test_validate_endpoints_sane_rate_and_burst() {
582        let endpoints = vec![
583            EndpointConfig {
584                name: None,
585                url: "https://good.com".into(),
586                rate_limit_per_sec: 0,
587                burst_size: 0,
588                weight: None,
589            },
590            EndpointConfig {
591                name: None,
592                url: "https://ok.com".into(),
593                rate_limit_per_sec: 1_000_000,
594                burst_size: 1_000_000,
595                weight: None,
596            },
597        ];
598        let result = validate_and_dedupe_endpoints(endpoints);
599        assert!(result.is_err(), "Endpoints with insane rate or burst should be rejected");
600    }
601
602    #[test]
603    fn test_validate_endpoints_canonicalization() {
604        let endpoints = vec![
605            EndpointConfig {
606                name: None,
607                url: "https://example.com".into(),
608                rate_limit_per_sec: 10,
609                burst_size: 5,
610                weight: None,
611            },
612            EndpointConfig {
613                name: None,
614                url: "https://example.com/".into(),
615                rate_limit_per_sec: 10,
616                burst_size: 5,
617                weight: None,
618            },
619            EndpointConfig {
620                name: None,
621                url: "HTTPS://example.com".into(),
622                rate_limit_per_sec: 10,
623                burst_size: 5,
624                weight: None,
625            },
626        ];
627        let result = validate_and_dedupe_endpoints(endpoints).unwrap();
628        assert_eq!(
629            result.len(),
630            1,
631            "Duplicates differing only by slash or case should be deduplicated"
632        );
633    }
634
635    // --- NEW TESTS ADDED BELOW ---
636
637    #[test]
638    fn test_validate_endpoints_invalid_weight() {
639        let endpoints = vec![EndpointConfig {
640            name: None,
641            url: "https://example.com".into(),
642            rate_limit_per_sec: 10,
643            burst_size: 5,
644            weight: Some(0),
645        }];
646        let result = validate_and_dedupe_endpoints(endpoints);
647        assert!(result.is_ok());
648        let list = result.unwrap();
649        assert_eq!(list[0].weight, Some(0));
650    }
651
652    #[test]
653    fn test_validate_endpoints_multiple_trailing_slashes() {
654        let endpoints = vec![EndpointConfig {
655            name: None,
656            url: "https://example.com///".into(),
657            rate_limit_per_sec: 10,
658            burst_size: 5,
659            weight: None,
660        }];
661        let result = validate_and_dedupe_endpoints(endpoints).unwrap();
662        assert_eq!(result[0].url, "https://example.com");
663    }
664
665    #[test]
666    fn test_validate_endpoints_error_message() {
667        let endpoints = vec![EndpointConfig {
668            name: None,
669            url: "invalid".into(),
670            rate_limit_per_sec: 10,
671            burst_size: 5,
672            weight: None,
673        }];
674        let result = validate_and_dedupe_endpoints(endpoints);
675        assert!(matches!(
676            result,
677            Err(ConfigError::ConfigError(msg)) if msg == "No valid endpoints configured"
678        ));
679    }
680
681    #[test]
682    fn test_balancer_config_client_settings() {
683        let toml = r#"
684        [balancer]
685        connect_timeout_ms = 1000
686        timeout_secs = 10
687        pool_idle_timeout_secs = 30
688        pool_max_idle_per_host = 100
689        "#;
690
691        let mut file = NamedTempFile::new().unwrap();
692        writeln!(file, "{}", toml).unwrap();
693        let path = file.path().to_str().unwrap();
694
695        let cfg = try_load_config(path).unwrap().unwrap();
696        let balancer = cfg.balancer.unwrap();
697
698        assert_eq!(balancer.connect_timeout_ms.unwrap(), 1000);
699        assert_eq!(balancer.timeout_secs.unwrap(), 10);
700        assert_eq!(balancer.pool_idle_timeout_secs.unwrap(), 30);
701        assert_eq!(balancer.pool_max_idle_per_host.unwrap(), 100);
702    }
703
704    #[test]
705    fn test_get_default_endpoints_urls() {
706        let endpoints = get_default_endpoints();
707        let expected_urls: Vec<String> = DEFAULT_ENDPOINTS.iter().map(|&s| s.to_string()).collect();
708        let actual_urls: Vec<String> = endpoints.iter().map(|e| e.url.clone()).collect();
709
710        assert_eq!(actual_urls, expected_urls);
711    }
712
713    #[test]
714    fn test_default_constants() {
715        let endpoint = EndpointConfig::default();
716        assert_eq!(endpoint.burst_size, DEFAULT_BURST_SIZE);
717        assert_eq!(endpoint.rate_limit_per_sec, DEFAULT_ENDPOINT_RATE_LIMIT);
718
719        let default_endpoints = get_default_endpoints();
720        assert!(default_endpoints.iter().all(|e| e.burst_size == DEFAULT_BURST_SIZE));
721        assert!(default_endpoints
722            .iter()
723            .all(|e| e.rate_limit_per_sec == DEFAULT_ENDPOINT_RATE_LIMIT));
724    }
725
726    #[test]
727    fn test_endpoint_scheme_mixed_case() {
728        let endpoints = vec![EndpointConfig {
729            name: None,
730            url: "HtTpS://example.com".into(),
731            rate_limit_per_sec: 10,
732            burst_size: 5,
733            weight: None,
734        }];
735        let validated = validate_and_dedupe_endpoints(endpoints).unwrap();
736        assert_eq!(validated[0].url, "https://example.com");
737    }
738
739    #[test]
740    fn test_endpoint_with_query_and_fragment() {
741        let endpoints = vec![EndpointConfig {
742            name: None,
743            url: "https://example.com/path?query=1#frag".into(),
744            rate_limit_per_sec: 10,
745            burst_size: 5,
746            weight: None,
747        }];
748        let validated = validate_and_dedupe_endpoints(endpoints).unwrap();
749        assert_eq!(validated[0].url, "https://example.com/path?query=1#frag");
750    }
751
752    #[test]
753    fn test_balancer_extreme_values() {
754        let cfg = BalancerConfig {
755            max_batch_size: Some(0),
756            max_concurrency: Some(100_000),
757            ..Default::default()
758        };
759        assert_eq!(cfg.max_batch_size.unwrap(), 0);
760        assert_eq!(cfg.max_concurrency.unwrap(), 100_000);
761    }
762
763    #[test]
764    fn test_latency_smoothing_edge() {
765        let cfg = BalancerConfig { latency_smoothing_factor: Some(0.0), ..Default::default() };
766        assert_eq!(cfg.latency_smoothing_factor.unwrap(), 0.0);
767    }
768
769    #[test]
770    fn test_validate_endpoints_all_invalid_branches() {
771        let endpoints = vec![
772            // Empty URL → invalid
773            EndpointConfig {
774                name: None,
775                url: "".into(),
776                rate_limit_per_sec: 10,
777                burst_size: 5,
778                weight: None,
779            },
780            // Invalid URL → invalid
781            EndpointConfig {
782                name: None,
783                url: "not-a-url".into(),
784                rate_limit_per_sec: 10,
785                burst_size: 5,
786                weight: None,
787            },
788            // Rate limit zero → invalid
789            EndpointConfig {
790                name: None,
791                url: "https://ok.com".into(),
792                rate_limit_per_sec: 0,
793                burst_size: 5,
794                weight: None,
795            },
796            // Burst size zero → invalid
797            EndpointConfig {
798                name: None,
799                url: "https://ok2.com".into(),
800                rate_limit_per_sec: 10,
801                burst_size: 0,
802                weight: None,
803            },
804            // URL too long
805            EndpointConfig {
806                name: None,
807                url: format!("https://example.com/{}", "a".repeat(2048)),
808                rate_limit_per_sec: 10,
809                burst_size: 5,
810                weight: None,
811            },
812        ];
813
814        let result = validate_and_dedupe_endpoints(endpoints);
815        assert!(matches!(
816            result,
817            Err(ConfigError::ConfigError(msg)) if msg == "No valid endpoints configured"
818        ));
819    }
820
821    #[test]
822    fn test_validate_endpoints_mixed_stress() {
823        let endpoints = vec![
824            // Valid endpoints
825            EndpointConfig {
826                name: None,
827                url: "https://valid1.com".into(),
828                rate_limit_per_sec: 10,
829                burst_size: 5,
830                weight: Some(1),
831            },
832            EndpointConfig {
833                name: None,
834                url: "http://valid2.com/".into(),
835                rate_limit_per_sec: 20,
836                burst_size: 10,
837                weight: None,
838            },
839            // Invalid endpoints
840            EndpointConfig {
841                name: None,
842                url: "".into(),
843                rate_limit_per_sec: 10,
844                burst_size: 5,
845                weight: None,
846            },
847            EndpointConfig {
848                name: None,
849                url: "bad-url".into(),
850                rate_limit_per_sec: 10,
851                burst_size: 5,
852                weight: None,
853            },
854            EndpointConfig {
855                name: None,
856                url: "https://example.com".into(),
857                rate_limit_per_sec: 0,
858                burst_size: 5,
859                weight: None,
860            },
861            EndpointConfig {
862                name: None,
863                url: "https://example.com".into(),
864                rate_limit_per_sec: 10,
865                burst_size: 0,
866                weight: None,
867            },
868            EndpointConfig {
869                name: None,
870                url: format!("https://toolong.com/{}", "a".repeat(2048)),
871                rate_limit_per_sec: 10,
872                burst_size: 5,
873                weight: None,
874            },
875        ];
876
877        let result = validate_and_dedupe_endpoints(endpoints).unwrap();
878
879        // Only the valid endpoints remain
880        assert_eq!(result.len(), 2);
881        assert!(result.iter().any(|e| e.url == "https://valid1.com"));
882        assert!(result.iter().any(|e| e.url == "http://valid2.com"));
883    }
884
885    #[test]
886    fn test_finalize_applies_defaults() {
887        let cfg = Config::default();
888        let finalized = cfg.finalize().unwrap();
889
890        assert_eq!(finalized.server.unwrap().bind_addr.unwrap(), DEFAULT_BIND_ADDR);
891        assert_eq!(
892            finalized.balancer.as_ref().unwrap().health_check_interval_secs.unwrap(),
893            DEFAULT_HEALTH_CHECK_INTERVAL_SECS
894        );
895        assert_eq!(
896            finalized.balancer.as_ref().unwrap().max_batch_size.unwrap(),
897            DEFAULT_MAX_BATCH_SIZE
898        );
899    }
900
901    #[test]
902    fn test_finalize_clamps_latency_smoothing() {
903        let cfg = Config {
904            balancer: Some(BalancerConfig {
905                latency_smoothing_factor: Some(1.5),
906                ..Default::default()
907            }),
908            ..Default::default()
909        };
910
911        let finalized = cfg.finalize().unwrap();
912        assert_eq!(finalized.balancer.unwrap().latency_smoothing_factor.unwrap(), 1.0);
913    }
914
915    #[test]
916    fn test_finalize_enforces_min_batch_size() {
917        let cfg = Config {
918            balancer: Some(BalancerConfig { max_batch_size: Some(0), ..Default::default() }),
919            ..Default::default()
920        };
921
922        let finalized = cfg.finalize().unwrap();
923        assert_eq!(finalized.balancer.unwrap().max_batch_size.unwrap(), 1);
924    }
925
926    #[test]
927    fn test_finalize_applies_default_weight() {
928        let cfg = Config {
929            balancer: Some(BalancerConfig {
930                endpoints: Some(vec![EndpointConfig {
931                    name: None,
932                    url: "https://test.com".into(),
933                    rate_limit_per_sec: 10,
934                    burst_size: 5,
935                    weight: None,
936                }]),
937                ..Default::default()
938            }),
939            ..Default::default()
940        };
941
942        let finalized = cfg.finalize().unwrap();
943        let endpoints = finalized.balancer.unwrap().endpoints.unwrap();
944        assert_eq!(endpoints[0].weight.unwrap(), DEFAULT_ENDPOINT_WEIGHT);
945    }
946
947    #[test]
948    fn test_resolve_env_vars_in_urls() {
949        std::env::set_var("TEST_ENDPOINT", "https://resolved.com");
950
951        let cfg = Config {
952            balancer: Some(BalancerConfig {
953                endpoints: Some(vec![EndpointConfig {
954                    name: None,
955                    url: "${TEST_ENDPOINT}".into(),
956                    rate_limit_per_sec: 10,
957                    burst_size: 5,
958                    weight: None,
959                }]),
960                ..Default::default()
961            }),
962            ..Default::default()
963        };
964
965        let finalized = cfg.finalize().unwrap();
966        let endpoints = finalized.balancer.unwrap().endpoints.unwrap();
967        assert_eq!(endpoints[0].url, "https://resolved.com");
968
969        std::env::remove_var("TEST_ENDPOINT");
970    }
971
972    #[test]
973    fn test_endpoint_name_field() {
974        let endpoint = EndpointConfig {
975            name: Some("primary".to_string()),
976            url: "https://test.com".into(),
977            rate_limit_per_sec: 10,
978            burst_size: 5,
979            weight: None,
980        };
981
982        assert_eq!(endpoint.name.unwrap(), "primary");
983    }
984
985    #[test]
986    fn test_finalize_with_empty_config() {
987        let cfg = Config::default();
988        let finalized = cfg.finalize().unwrap();
989
990        // Should have default endpoints
991        let endpoints = finalized.balancer.unwrap().endpoints.unwrap();
992        assert!(!endpoints.is_empty());
993    }
994
995    #[test]
996    fn test_finalize_enforces_min_concurrency() {
997        let cfg = Config {
998            balancer: Some(BalancerConfig { max_concurrency: Some(0), ..Default::default() }),
999            ..Default::default()
1000        };
1001
1002        let finalized = cfg.finalize().unwrap();
1003        assert_eq!(finalized.balancer.unwrap().max_concurrency.unwrap(), 1);
1004    }
1005
1006    #[test]
1007    fn test_resolve_env_vars_missing() {
1008        // Use a valid base URL with env var substitution
1009        let cfg = Config {
1010            balancer: Some(BalancerConfig {
1011                endpoints: Some(vec![EndpointConfig {
1012                    name: None,
1013                    url: "https://${MISSING_VAR}.example.com".into(), // Valid URL format with env var
1014                    rate_limit_per_sec: 10,
1015                    burst_size: 5,
1016                    weight: None,
1017                }]),
1018                ..Default::default()
1019            }),
1020            ..Default::default()
1021        };
1022
1023        let finalized = cfg.finalize().unwrap();
1024        let endpoints = finalized.balancer.unwrap().endpoints.unwrap();
1025        // Should keep the original value if env var not found
1026        assert_eq!(endpoints[0].url, "https://${MISSING_VAR}.example.com");
1027    }
1028
1029    #[test]
1030    fn test_finalize_enforces_min_rate_limit() {
1031        // Create an endpoint with valid URL but rate limit that will be clamped
1032        // Note: rate_limit_per_sec = 1 is valid, but let's test the clamping logic
1033        let cfg = Config {
1034            balancer: Some(BalancerConfig {
1035                endpoints: Some(vec![EndpointConfig {
1036                    name: None,
1037                    url: "https://test.com".into(),
1038                    rate_limit_per_sec: 1, // Valid but let's see if it gets clamped
1039                    burst_size: 5,
1040                    weight: None,
1041                }]),
1042                ..Default::default()
1043            }),
1044            ..Default::default()
1045        };
1046
1047        let finalized = cfg.finalize().unwrap();
1048        let endpoints = finalized.balancer.unwrap().endpoints.unwrap();
1049        // Should remain 1 (no clamping needed since it's already >= 1)
1050        assert_eq!(endpoints[0].rate_limit_per_sec, 1);
1051    }
1052
1053    // Add this new test to specifically test the clamping of very low but valid rate limits
1054    #[test]
1055    fn test_finalize_clamps_low_but_valid_rate_limit() {
1056        let cfg = Config {
1057            balancer: Some(BalancerConfig {
1058                endpoints: Some(vec![EndpointConfig {
1059                    name: None,
1060                    url: "https://test.com".into(),
1061                    rate_limit_per_sec: 1, // Valid minimum
1062                    burst_size: 5,
1063                    weight: None,
1064                }]),
1065                ..Default::default()
1066            }),
1067            ..Default::default()
1068        };
1069
1070        let finalized = cfg.finalize().unwrap();
1071        let endpoints = finalized.balancer.unwrap().endpoints.unwrap();
1072        // Should remain 1 (no change needed)
1073        assert_eq!(endpoints[0].rate_limit_per_sec, 1);
1074    }
1075}
azoth_balancer/config.rs

azoth_balancer/
config.rs
