Skip to content

Commit

Permalink
small change
Browse files Browse the repository at this point in the history
  • Loading branch information
hatoo committed Dec 22, 2024
1 parent 6cb4caa commit c084e3c
Show file tree
Hide file tree
Showing 2 changed files with 46 additions and 42 deletions.
21 changes: 19 additions & 2 deletions src/main.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;

Expand Down Expand Up @@ -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<Url> = reader
.lines()
.map_while(Result::ok)
.filter(|line| !line.trim().is_empty())
.map(|url_str| Url::parse(&url_str))
.collect::<Result<Vec<_>, _>>()?;
UrlGenerator::new_multi_static(urls)
} else {
UrlGenerator::new_static(Url::parse(&opts.url)?)
};
Expand Down
67 changes: 27 additions & 40 deletions src/url_generator.rs
Original file line number Diff line number Diff line change
@@ -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::*;
Expand Down Expand Up @@ -33,23 +30,9 @@ impl UrlGenerator {
Self::Static(url)
}

pub fn new_multi_static(filename: &str) -> Result<Self, UrlGeneratorError> {
let path = Path::new(filename);
let file = File::open(path)?;
let reader = io::BufReader::new(file);

let urls: Vec<Url> = reader
.lines()
.map_while(Result::ok)
.filter(|line| !line.trim().is_empty())
.map(|url_str| {
Url::parse(&url_str).map_err(|e| {
UrlGeneratorError::Parse(e, format!("Failed to parse URL '{}': {}", url_str, e))
})
})
.collect::<Result<Vec<_>, _>>()?;

Ok(Self::MultiStatic(urls))
pub fn new_multi_static(urls: Vec<Url>) -> Self {
assert!(!urls.is_empty());
Self::MultiStatic(urls)
}

pub fn new_dynamic(regex: Regex) -> Self {
Expand Down Expand Up @@ -84,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]
Expand All @@ -98,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![
"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();
Expand All @@ -117,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]";
Expand Down Expand Up @@ -160,4 +125,26 @@ mod tests {
);
}
}

#[test]
fn test_url_generator_multi_consistency() {
let urls = vec![
"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()
);
}
}
}

0 comments on commit c084e3c

Please sign in to comment.