1use std::{
2 fs,
3 io::{self, BufRead},
4 path,
5};
6
7use toml::Value;
8
9#[derive(thiserror::Error, Debug)]
10pub enum ConfigError {
11 #[error("IO error: {0}")]
12 Io(#[from] std::io::Error),
13 #[error("TOML parsing error: {0}")]
14 TomlParse(#[from] toml::de::Error),
15 #[error("Line parsing error: {0}")]
16 LineParse(String),
17}
18
19pub trait Configurable {
20 fn config(&self) -> &Value;
21
22 fn load_config(
24 config_file_path: impl AsRef<path::Path>,
25 ) -> Result<Value, ConfigError> {
26 let content: String = fs::read_to_string(config_file_path)?;
27 let config: Value = toml::from_str(&content)?;
28 Ok(config)
29 }
30
31 fn load_text_file_lines(
33 file_path: impl AsRef<path::Path>,
34 ) -> Result<Vec<String>, ConfigError> {
35 let file = fs::File::open(file_path)?;
36 let lines = io::BufReader::new(file)
37 .lines()
38 .map(|l| l.map_err(|e| ConfigError::LineParse(e.to_string())))
39 .collect::<Result<Vec<_>, _>>()?;
40 Ok(lines)
41 }
42
43 fn get_config_value(&self, key: &str) -> Option<&Value> {
45 let keys: Vec<&str> = key.split('.').collect();
46 Self::get_value_recursive(self.config(), &keys)
47 }
48
49 fn get_value_recursive<'a>(
50 config: &'a Value,
51 keys: &[&str],
52 ) -> Option<&'a Value> {
53 if keys.is_empty() {
54 return None;
55 };
56
57 match config {
58 Value::Table(map) => {
59 let key = keys[0];
60 let remaining_keys = &keys[1..];
61
62 if let Some(value) = map.get(key) {
63 if remaining_keys.is_empty() {
64 Some(value)
65 } else {
66 Self::get_value_recursive(value, remaining_keys)
67 }
68 } else {
69 None
70 }
71 }
72 _ => None,
73 }
74 }
75}
76
77#[cfg(test)]
78mod tests {
79 use super::*;
80 use std::fs::File;
81 use std::io::Write;
82 use tempfile::tempdir;
83
84 pub struct TestApp {
85 config: Value,
86 user_agents: Option<Vec<String>>,
87 }
88
89 impl Configurable for TestApp {
90 fn config(&self) -> &Value {
91 &self.config
92 }
93 }
94
95 impl TestApp {
96 fn from_config(config_file_path: impl AsRef<path::Path>) -> Self {
97 let config = Self::load_config(config_file_path);
98 Self {
99 config: config.unwrap(),
100 user_agents: None,
101 }
102 }
103
104 fn load_uas(&mut self, uas_file_path: impl AsRef<path::Path>) {
105 self.user_agents = Self::load_text_file_lines(uas_file_path).ok();
106 }
107 }
108
109 #[test]
110 fn test_load_config() {
111 let config_path = "../tests/simple_config.toml";
112 let app = TestApp::from_config(config_path);
113
114 assert_eq!(app.config["app"]["threads"].as_integer(), Some(4));
115 assert_eq!(app.config()["app"]["max_queue"].as_integer(), Some(500));
116 assert_eq!(app.user_agents, None);
117 }
118
119 #[test]
120 fn test_load_config_valid_toml() {
121 let dir = tempdir().unwrap();
122 let config_path = dir.path().join("config.toml");
123 let mut file = File::create(&config_path).unwrap();
124 writeln!(file, "key = \"value\"\n[app]\nsetting = 42").unwrap();
125
126 let config = TestApp::load_config(&config_path);
127 assert!(config.is_ok());
128 let config = config.unwrap();
129 assert_eq!(config["key"].as_str(), Some("value"));
130 assert_eq!(config["app"]["setting"].as_integer(), Some(42));
131 }
132
133 #[test]
134 fn test_load_config_invalid_toml() {
135 let dir = tempdir().unwrap();
136 let config_path = dir.path().join("config.toml");
137 let mut file = File::create(&config_path).unwrap();
138 writeln!(file, "invalid = : toml").unwrap();
139
140 let config = TestApp::load_config(&config_path);
141 assert!(matches!(config, Err(ConfigError::TomlParse(_))));
142 }
143
144 #[test]
145 fn test_load_text_file_lines() {
146 let dir = tempdir().unwrap();
147 let file_path = dir.path().join("test.txt");
148 let mut file = File::create(&file_path).unwrap();
149 writeln!(file, "line1\nline2\nline3").unwrap();
150
151 let lines = TestApp::load_text_file_lines(&file_path);
152 assert!(lines.is_ok());
153 let lines = lines.unwrap();
154 assert_eq!(lines, vec!["line1", "line2", "line3"]);
155 }
156
157 #[test]
158 fn test_get_config_value_empty_keys() {
159 let config_path = "../tests/simple_config.toml";
160 let app = TestApp::from_config(config_path);
161 assert_eq!(app.get_config_value(""), None);
162 }
163
164 #[test]
165 fn test_get_config_value() {
166 let config_path = "../tests/simple_config.toml";
167 let app = TestApp::from_config(config_path);
168
169 assert_eq!(
170 app.get_config_value("logging.log_to_file")
171 .unwrap()
172 .as_bool(),
173 Some(true)
174 )
175 }
176
177 #[test]
178 fn test_get_config_value_recursive() {
179 let toml = r#"
180 [app.nested]
181 value = 42
182 "#;
183 let config: Value = toml::from_str(toml).unwrap();
184 let app = TestApp {
185 config,
186 user_agents: None,
187 };
188
189 assert_eq!(
190 app.get_config_value("app.nested.value")
191 .and_then(|v| v.as_integer()),
192 Some(42)
193 );
194 assert_eq!(app.get_config_value("app.missing.value"), None);
195 assert_eq!(app.get_config_value("missing"), None);
196 }
197
198 #[test]
199 fn test_load_lines() {
200 let config_path = "../tests/simple_config.toml";
201 let mut app = TestApp::from_config(config_path);
202 let uas_file_path = {
203 app.get_config_value("app.user_agents_file")
204 .unwrap()
205 .as_str()
206 .unwrap()
207 .to_owned()
208 };
209 app.load_uas(&uas_file_path);
210 }
211}