diff --git a/src/main.rs b/src/main.rs index 36ffca83..3261e71b 100644 --- a/src/main.rs +++ b/src/main.rs @@ -12,7 +12,14 @@ use rand::prelude::*; use rand_regex::Regex; use ratatui::crossterm; use result_data::ResultData; -use std::{env, io::Read, str::FromStr, sync::Arc}; +use std::{ + env, + fs::File, + io::{BufRead, Read}, + path::Path, + str::FromStr, + sync::Arc, +}; use url::Url; use url_generator::UrlGenerator; @@ -327,7 +334,17 @@ async fn main() -> anyhow::Result<()> { .collect(); UrlGenerator::new_dynamic(Regex::compile(&dot_disabled, opts.max_repeat)?) } else if opts.urls_from_file { - UrlGenerator::new_multi_static(&opts.url)? + let path = Path::new(opts.url.as_str()); + let file = File::open(path)?; + let reader = std::io::BufReader::new(file); + + let urls: Vec = reader + .lines() + .map_while(Result::ok) + .filter(|line| !line.trim().is_empty()) + .map(|url_str| Url::parse(&url_str)) + .collect::, _>>()?; + UrlGenerator::new_multi_static(urls) } else { UrlGenerator::new_static(Url::parse(&opts.url)?) }; diff --git a/src/url_generator.rs b/src/url_generator.rs index 1f31bf3e..0d56b510 100644 --- a/src/url_generator.rs +++ b/src/url_generator.rs @@ -1,6 +1,3 @@ -use std::fs::File; -use std::io::{self, BufRead}; -use std::path::Path; use std::{borrow::Cow, string::FromUtf8Error}; use rand::prelude::*; @@ -19,13 +16,13 @@ pub enum UrlGenerator { #[derive(Error, Debug)] pub enum UrlGeneratorError { #[error("{0}, generated url: {1}")] - ParseError(ParseError, String), + Parse(ParseError, String), #[error(transparent)] - FromUtf8Error(#[from] FromUtf8Error), + FromUtf8(#[from] FromUtf8Error), #[error("No valid URLs found")] - NoURLsError(), + NoURLs(), #[error(transparent)] - IoError(#[from] std::io::Error), + Io(#[from] std::io::Error), } impl UrlGenerator { @@ -33,26 +30,9 @@ impl UrlGenerator { Self::Static(url) } - pub fn new_multi_static(filename: &str) -> Result { - let path = Path::new(filename); - let file = File::open(path)?; - let reader = io::BufReader::new(file); - - let urls: Vec = reader - .lines() - .map_while(Result::ok) - .filter(|line| !line.trim().is_empty()) - .map(|url_str| { - Url::parse(&url_str).map_err(|e| { - UrlGeneratorError::ParseError( - e, - format!("Failed to parse URL '{}': {}", url_str, e), - ) - }) - }) - .collect::, _>>()?; - - Ok(Self::MultiStatic(urls)) + pub fn new_multi_static(urls: Vec) -> Self { + assert!(!urls.is_empty()); + Self::MultiStatic(urls) } pub fn new_dynamic(regex: Regex) -> Self { @@ -66,14 +46,15 @@ impl UrlGenerator { if let Some(random_url) = urls.choose(rng) { Ok(Cow::Borrowed(random_url)) } else { - Err(UrlGeneratorError::NoURLsError()) + Err(UrlGeneratorError::NoURLs()) } } Self::Dynamic(regex) => { let generated = Distribution::>::sample(regex, rng)?; - Ok(Cow::Owned(Url::parse(generated.as_str()).map_err(|e| { - UrlGeneratorError::ParseError(e, generated) - })?)) + Ok(Cow::Owned( + Url::parse(generated.as_str()) + .map_err(|e| UrlGeneratorError::Parse(e, generated))?, + )) } } } @@ -86,8 +67,7 @@ mod tests { use super::*; use rand_regex::Regex as RandRegex; use regex::Regex; - use std::{fs, net::Ipv4Addr}; - use tempfile::NamedTempFile; + use std::net::Ipv4Addr; use url::{Host, Url}; #[test] @@ -100,17 +80,14 @@ mod tests { #[test] fn test_url_generator_multistatic() { - let temp_file = NamedTempFile::new().expect("Failed to create temporary file"); - let urls = vec![ + let urls = [ "http://127.0.0.1/a1", "http://127.0.0.1/b2", "http://127.0.0.1/c3", ]; - let file_content = urls.join("\n") + "\n\n"; - fs::write(temp_file.path(), file_content).expect("Failed to write to temporary file"); let url_generator = - UrlGenerator::new_multi_static(temp_file.path().to_str().unwrap()).unwrap(); + UrlGenerator::new_multi_static(urls.iter().map(|u| Url::parse(u).unwrap()).collect()); for _ in 0..10 { let url = url_generator.generate(&mut thread_rng()).unwrap(); @@ -119,20 +96,6 @@ mod tests { } } - #[test] - fn test_url_generator_multistatic_errors() { - let temp_file = NamedTempFile::new().expect("Failed to create temporary file"); - let file_content = "https://127.0.0.1\n\nno url\n"; - fs::write(temp_file.path(), file_content).expect("Failed to write to temporary file"); - - let url_generator = UrlGenerator::new_multi_static(temp_file.path().to_str().unwrap()); - - assert!( - url_generator.is_err(), - "Parsing should have failed with an error!" - ); - } - #[test] fn test_url_generator_dynamic() { let path_regex = "/[a-z][a-z][0-9]"; @@ -162,4 +125,26 @@ mod tests { ); } } + + #[test] + fn test_url_generator_multi_consistency() { + let urls = [ + "http://example.com/a1", + "http://example.com/a2", + "http://example.com/a3", + "http://example.com/a4", + "http://example.com/a5", + ]; + let url_generator = + UrlGenerator::new_multi_static(urls.iter().map(|u| Url::parse(u).unwrap()).collect()); + + for _ in 0..100 { + let rng: Pcg64Si = SeedableRng::from_entropy(); + + assert_eq!( + url_generator.generate(&mut rng.clone()).unwrap(), + url_generator.generate(&mut rng.clone()).unwrap() + ); + } + } }