diff --git a/crates/irtt-cli/Cargo.toml b/crates/irtt-cli/Cargo.toml index d77094a..6e490bf 100644 --- a/crates/irtt-cli/Cargo.toml +++ b/crates/irtt-cli/Cargo.toml @@ -55,3 +55,7 @@ irtt-server = { workspace = true, optional = true } irtt-stats = { workspace = true, optional = true } ratatui = { version = "0.30.2", optional = true } tokio = { version = "1.53.1", default-features = false, features = ["rt", "time"], optional = true } + +[dev-dependencies] +irtt-server.workspace = true +tokio = { version = "1.53.1", default-features = false, features = ["rt", "sync"] } diff --git a/crates/irtt-cli/tests/client_multi_target.rs b/crates/irtt-cli/tests/client_multi_target.rs index c8ab20d..0f2cc93 100644 --- a/crates/irtt-cli/tests/client_multi_target.rs +++ b/crates/irtt-cli/tests/client_multi_target.rs @@ -1,5 +1,8 @@ #![cfg(feature = "client")] +#[path = "support/in_tree_server.rs"] +mod in_tree_server; + use std::{ io::Read, net::{SocketAddr, UdpSocket}, @@ -10,6 +13,7 @@ use std::{ time::{Duration, Instant}, }; +use in_tree_server::InTreeServer; use irtt_proto::{ echo_packet_len, encode_echo_reply, encode_open_reply, flags::{FLAG_CLOSE, FLAG_OPEN, FLAG_REPLY}, @@ -69,10 +73,7 @@ fn list_columns_succeeds_without_target() { #[test] fn single_target_default_table_omits_target_and_accepts_pacing() { - let server = start_echo_server(test_params( - Some(Duration::from_millis(30)), - Duration::from_millis(10), - )); + let server = InTreeServer::start(); let output = Command::new(env!("CARGO_BIN_EXE_irtt-cli")) .args([ @@ -89,7 +90,7 @@ fn single_target_default_table_omits_target_and_accepts_pacing() { .output() .unwrap(); - server.join(); + drop(server); assert!( output.status.success(), @@ -104,10 +105,7 @@ fn single_target_default_table_omits_target_and_accepts_pacing() { #[test] fn single_target_custom_target_column_renders_positional_label() { - let server = start_echo_server(test_params( - Some(Duration::from_millis(30)), - Duration::from_millis(10), - )); + let server = InTreeServer::start(); let target = server.addr.to_string(); let output = Command::new(env!("CARGO_BIN_EXE_irtt-cli")) @@ -127,7 +125,7 @@ fn single_target_custom_target_column_renders_positional_label() { .output() .unwrap(); - server.join(); + drop(server); assert!( output.status.success(), @@ -238,9 +236,8 @@ fn continuous_single_target_interruption_drains_final_events_and_succeeds() { #[test] fn multi_target_csv_emits_rows_for_both_labels() { - let params = test_params(Some(Duration::from_millis(40)), Duration::from_millis(10)); - let a = start_echo_server(params.clone()); - let b = start_echo_server(params); + let a = InTreeServer::start(); + let b = InTreeServer::start(); let output = Command::new(env!("CARGO_BIN_EXE_irtt-cli")) .args([ @@ -262,8 +259,8 @@ fn multi_target_csv_emits_rows_for_both_labels() { .output() .unwrap(); - a.join(); - b.join(); + drop(a); + drop(b); assert!( output.status.success(), @@ -643,20 +640,9 @@ fn wait_with_timeout(mut child: std::process::Child, timeout: Duration) -> Outpu } } -fn start_echo_server(params: Params) -> FakeServer { - start_echo_server_inner(params, None) -} - fn start_echo_server_with_first_reply( params: Params, first_reply_tx: mpsc::Sender<()>, -) -> FakeServer { - start_echo_server_inner(params, Some(first_reply_tx)) -} - -fn start_echo_server_inner( - params: Params, - mut first_reply_tx: Option>, ) -> FakeServer { let socket = UdpSocket::bind("127.0.0.1:0").unwrap(); socket @@ -664,6 +650,7 @@ fn start_echo_server_inner( .unwrap(); let addr = socket.local_addr().unwrap(); let done = thread::spawn(move || { + let mut first_reply_tx = Some(first_reply_tx); let (_, peer) = recv_request(&socket); socket .send_to(&open_reply(FLAG_OPEN | FLAG_REPLY, TOKEN, ¶ms), peer) diff --git a/crates/irtt-cli/tests/support/in_tree_server.rs b/crates/irtt-cli/tests/support/in_tree_server.rs new file mode 100644 index 0000000..767731f --- /dev/null +++ b/crates/irtt-cli/tests/support/in_tree_server.rs @@ -0,0 +1,74 @@ +use std::{ + net::{IpAddr, Ipv4Addr, SocketAddr}, + sync::mpsc, + thread::{self, JoinHandle}, + time::Duration, +}; + +use irtt_server::{Server, ServerConfig}; +use tokio::sync::oneshot; + +const STARTUP_TIMEOUT: Duration = Duration::from_secs(2); + +pub struct InTreeServer { + pub addr: SocketAddr, + shutdown: Option>, + done: Option>>, +} + +impl InTreeServer { + pub fn start() -> Self { + let (startup_tx, startup_rx) = mpsc::sync_channel(1); + let (shutdown_tx, shutdown_rx) = oneshot::channel(); + let done = thread::spawn(move || { + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .expect("failed to build in-tree test server runtime"); + runtime.block_on(async move { + let bind_addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 0); + let mut server = Server::bind(bind_addr, ServerConfig::default()).await?; + startup_tx + .send(server.local_addr()?) + .expect("test stopped waiting for in-tree server startup"); + server + .run(async { + let _ = shutdown_rx.await; + }) + .await + }) + }); + + let addr = match startup_rx.recv_timeout(STARTUP_TIMEOUT) { + Ok(addr) => addr, + Err(error) => { + let _ = shutdown_tx.send(()); + let result = done.join().expect("in-tree test server thread panicked"); + panic!("in-tree test server failed before startup ({error}): {result:?}"); + } + }; + + Self { + addr, + shutdown: Some(shutdown_tx), + done: Some(done), + } + } + + fn stop(&mut self) { + if let Some(shutdown) = self.shutdown.take() { + let _ = shutdown.send(()); + } + if let Some(done) = self.done.take() { + done.join() + .expect("in-tree test server thread panicked") + .expect("in-tree test server stopped with an error"); + } + } +} + +impl Drop for InTreeServer { + fn drop(&mut self) { + self.stop(); + } +} diff --git a/crates/irtt-client/tests/support/in_tree_server.rs b/crates/irtt-client/tests/support/in_tree_server.rs index 95680a0..faf4fcd 100644 --- a/crates/irtt-client/tests/support/in_tree_server.rs +++ b/crates/irtt-client/tests/support/in_tree_server.rs @@ -13,7 +13,7 @@ const STARTUP_TIMEOUT: Duration = Duration::from_secs(2); pub struct InTreeServer { pub addr: SocketAddr, shutdown: Option>, - done: Option>>, + done: Option>>, } impl InTreeServer { @@ -24,52 +24,27 @@ impl InTreeServer { let runtime = tokio::runtime::Builder::new_current_thread() .enable_all() .build() - .map_err(|error| format!("failed to build test server runtime: {error}"))?; + .expect("failed to build in-tree test server runtime"); runtime.block_on(async move { let bind_addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 0); - let mut server = match Server::bind(bind_addr, config).await { - Ok(server) => server, - Err(error) => { - let message = format!("failed to bind in-tree test server: {error}"); - let _ = startup_tx.send(Err(message.clone())); - return Err(message); - } - }; - let addr = match server.local_addr() { - Ok(addr) => addr, - Err(error) => { - let message = - format!("failed to read in-tree test server address: {error}"); - let _ = startup_tx.send(Err(message.clone())); - return Err(message); - } - }; + let mut server = Server::bind(bind_addr, config).await?; startup_tx - .send(Ok(addr)) - .map_err(|_| "test stopped waiting for in-tree server startup".to_owned())?; + .send(server.local_addr()?) + .expect("test stopped waiting for in-tree server startup"); server .run(async { let _ = shutdown_rx.await; }) .await - .map_err(|error| format!("in-tree test server failed: {error}")) }) }); let addr = match startup_rx.recv_timeout(STARTUP_TIMEOUT) { - Ok(Ok(addr)) => addr, - Ok(Err(error)) => { - done.join() - .expect("in-tree test server thread panicked") - .ok(); - panic!("{error}"); - } + Ok(addr) => addr, Err(error) => { let _ = shutdown_tx.send(()); - done.join() - .expect("in-tree test server thread panicked") - .ok(); - panic!("timed out waiting for in-tree test server startup: {error}"); + let result = done.join().expect("in-tree test server thread panicked"); + panic!("in-tree test server failed before startup ({error}): {result:?}"); } };