diff --git a/ubw-sward/src/http/random.rs b/ubw-sward/src/http/random.rs index 8274183..f1a87af 100644 --- a/ubw-sward/src/http/random.rs +++ b/ubw-sward/src/http/random.rs @@ -4,12 +4,20 @@ use regex::Regex; use url::Url; #[derive(Clone)] +/// A Regex-based URL generator. pub struct RandomUrlGenerator { random_engine: Rng, random_parts: Box<[RegexPart]>, } impl RandomUrlGenerator { + /// Create a new RandomUrlGenerator. + /// + /// # Arguments + /// + /// * `random_engine` - The random engine to use. + /// * `template` - The template to use. + /// pub fn new(random_engine: Rng, template: &str) -> Result { let re = Regex::new(r"\[([^\]]+)\](?:\{(\d+)\})?")?; let mut parts: Vec = Vec::new(); @@ -71,6 +79,7 @@ impl RandomUrlGenerator { }) } + /// Generate a random URL (string) pub fn generate(&mut self) -> CompactString { self.random_parts.iter().fold( CompactString::with_capacity(self.expected_length()), @@ -91,6 +100,7 @@ impl RandomUrlGenerator { ) } + /// Generate a random URL. Will return error if the URL is invalid. pub fn generate_url(&mut self) -> Result { let url = self.generate(); Url::parse(&url) diff --git a/ubw-sward/src/http/simple.rs b/ubw-sward/src/http/simple.rs index b0938a2..e2ecfab 100644 --- a/ubw-sward/src/http/simple.rs +++ b/ubw-sward/src/http/simple.rs @@ -2,6 +2,7 @@ use bytes::Bytes; use reqwest::header::HeaderMap; use reqwest::{Client, Method}; use std::pin::Pin; +use std::sync::Arc; use std::task::{Context, Poll}; use tower::Service; use url::Url; @@ -11,12 +12,13 @@ use url::Url; pub struct SimpleHttpSward { client: Client, method: Method, - headers: HeaderMap, + headers: Arc, sent_count: usize, } impl SimpleHttpSward { - pub fn new(client: Client, method: Method, headers: HeaderMap) -> Self { + /// Create a new simple http sward. + pub fn new(client: Client, method: Method, headers: Arc) -> Self { Self { client, method, @@ -24,6 +26,8 @@ impl SimpleHttpSward { sent_count: 0, } } + + /// Get the number of requests sent. pub fn sent_count(&self) -> usize { self.sent_count } @@ -47,7 +51,7 @@ impl Service for SimpleHttpSward { Box::pin( self.client .request(self.method.clone(), req.url) - .headers(self.headers.clone()) + .headers(self.headers.as_ref().clone()) .send(), ) }