diff --git a/lib/runtime/src/pipeline/network/tcp.rs b/lib/runtime/src/pipeline/network/tcp.rs index 9fda4c241..49e5d75df 100644 --- a/lib/runtime/src/pipeline/network/tcp.rs +++ b/lib/runtime/src/pipeline/network/tcp.rs @@ -20,6 +20,8 @@ pub mod client; pub mod server; +pub mod test_utils; + use super::ControlMessage; use serde::{Deserialize, Serialize}; diff --git a/lib/runtime/src/pipeline/network/tcp/client.rs b/lib/runtime/src/pipeline/network/tcp/client.rs index 2d6ba4577..ddb6b7738 100644 --- a/lib/runtime/src/pipeline/network/tcp/client.rs +++ b/lib/runtime/src/pipeline/network/tcp/client.rs @@ -324,3 +324,399 @@ async fn handle_writer( drop(alive_rx); Ok(framed_writer) } + +#[cfg(test)] +mod tests { + use super::*; + use crate::pipeline::context::Controller; + use crate::pipeline::network::tcp::test_utils::create_tcp_pair; + use bytes::Bytes; + use futures::StreamExt; + use std::sync::Arc; + use tokio::io::AsyncReadExt; + use tokio::net::TcpStream; + use tokio::sync::{mpsc, oneshot}; + use tokio_util::codec::FramedRead; + + struct WriterHarness { + server: tokio::net::TcpStream, + framed_writer: FramedWrite, TwoPartCodec>, + bytes_tx: mpsc::Sender, + bytes_rx: mpsc::Receiver, + alive_tx: oneshot::Sender<()>, + alive_rx: oneshot::Receiver<()>, + controller: Arc, + } + + /// Creates a reusable writer harness with paired TCP streams and test channels. + async fn writer_harness() -> WriterHarness { + let (client, server) = create_tcp_pair().await; + let (_, write_half) = tokio::io::split(client); + let framed_writer = FramedWrite::new(write_half, TwoPartCodec::default()); + + let (bytes_tx, bytes_rx) = mpsc::channel(64); + let (alive_tx, alive_rx) = oneshot::channel::<()>(); + let controller = Arc::new(Controller::default()); + + WriterHarness { + server, + framed_writer, + bytes_tx, + bytes_rx, + alive_tx, + alive_rx, + controller, + } + } + + async fn recv_msg(reader: &mut FramedRead) -> TwoPartMessage { + reader + .next() + .await + .expect("expected message") + .expect("failed to decode message") + } + + fn assert_data_only_message(msg: TwoPartMessage, expected: &[u8]) { + let (header, data) = msg.optional_parts(); + assert!(header.is_none(), "data-only message should not have header"); + assert_eq!( + data.expect("data payload missing").as_ref(), + expected, + "data payload should match" + ); + } + + fn assert_header_only_message(msg: TwoPartMessage, expected: &[u8]) { + let (header, data) = msg.optional_parts(); + assert!(data.is_none(), "header-only message should not carry data"); + assert_eq!( + header.expect("header missing").as_ref(), + expected, + "header payload should match" + ); + } + + fn assert_header_and_data_message( + msg: TwoPartMessage, + expected_header: &[u8], + expected_data: &[u8], + ) { + let (header, data) = msg.optional_parts(); + assert_eq!( + header.expect("header missing").as_ref(), + expected_header, + "header payload should match" + ); + assert_eq!( + data.expect("data missing").as_ref(), + expected_data, + "data payload should match" + ); + } + + fn assert_sentinel_message(msg: TwoPartMessage) { + let (header, data) = msg.optional_parts(); + assert!(data.is_none(), "sentinel should not include a data section"); + let expected_sentinel = serde_json::to_vec(&ControlMessage::Sentinel).unwrap(); + assert_eq!( + header.expect("sentinel header missing").as_ref(), + expected_sentinel.as_slice(), + "sentinel header should match serialized ControlMessage::Sentinel" + ); + } + + /// Test that handle_writer forwards messages from the channel to the framed writer + #[tokio::test] + async fn test_handle_writer_forwards_messages() { + let WriterHarness { + server, + framed_writer, + bytes_tx, + bytes_rx, + alive_rx, + controller, + .. + } = writer_harness().await; + + // Send test messages + let test_msg = TwoPartMessage::from_data(Bytes::from("test data")); + bytes_tx.send(test_msg).await.unwrap(); + + // Close the sender to trigger normal termination + drop(bytes_tx); + + let result = handle_writer(framed_writer, bytes_rx, alive_rx, controller).await; + + assert!(result.is_ok()); + + // Decode from server side to verify data and sentinel were sent + let mut reader = FramedRead::new(server, TwoPartCodec::default()); + + let msg = recv_msg(&mut reader).await; + assert_data_only_message(msg, b"test data"); + + let sentinel = recv_msg(&mut reader).await; + assert_sentinel_message(sentinel); + } + + /// Test that handle_writer sends sentinel on normal channel closure + #[tokio::test] + async fn test_handle_writer_sends_sentinel_on_normal_closure() { + let WriterHarness { + mut server, + framed_writer, + bytes_tx, + bytes_rx, + alive_rx, + controller, + .. + } = writer_harness().await; + + // Close the sender immediately to trigger normal termination + drop(bytes_tx); + + let result = handle_writer(framed_writer, bytes_rx, alive_rx, controller).await; + + assert!(result.is_ok()); + + // Read from server side to verify sentinel was sent + let mut buffer = vec![0u8; 1024]; + let n = server.read(&mut buffer).await.unwrap(); + + // Buffer should contain the sentinel message + assert!(n > 0, "Expected sentinel to be written to the TCP stream"); + + // Verify it contains the sentinel message by checking for the JSON + let sentinel_json = serde_json::to_vec(&ControlMessage::Sentinel).unwrap(); + assert!( + buffer[..n] + .windows(sentinel_json.len()) + .any(|w| w == sentinel_json.as_slice()), + "Buffer should contain sentinel message. Buffer: {:?}", + String::from_utf8_lossy(&buffer[..n]) + ); + } + + /// Test that handle_writer does NOT send sentinel when context is killed + #[tokio::test] + async fn test_handle_writer_no_sentinel_on_context_killed() { + let WriterHarness { + mut server, + framed_writer, + bytes_rx, + alive_rx, + controller, + .. + } = writer_harness().await; + + // Kill the context + controller.kill(); + + let result = handle_writer(framed_writer, bytes_rx, alive_rx, controller).await; + + assert!(result.is_ok()); + + // Drop the writer to close the connection, then try to read. Otherwise, + // the test will hang on `server.read()` + drop(result); + + // Read from server side - should get no sentinel + let mut buffer = vec![0u8; 1024]; + let n = server.read(&mut buffer).await.unwrap(); + + // Buffer should be empty (no sentinel sent) + let sentinel_json = serde_json::to_vec(&ControlMessage::Sentinel).unwrap(); + assert!( + n == 0 + || !buffer[..n] + .windows(sentinel_json.len()) + .any(|w| w == sentinel_json.as_slice()), + "Buffer should NOT contain sentinel message when context is killed" + ); + } + + /// Test that handle_writer does NOT send sentinel when context is stopped + #[tokio::test] + async fn test_handle_writer_no_sentinel_on_context_stopped() { + let WriterHarness { + mut server, + framed_writer, + bytes_rx, + alive_rx, + controller, + .. + } = writer_harness().await; + + // Stop the context + controller.stop(); + + let result = handle_writer(framed_writer, bytes_rx, alive_rx, controller).await; + + assert!(result.is_ok()); + + // Drop the writer to close the connection, then try to read. Otherwise, + // the test will hang on `server.read()` + drop(result); + + // Read from server side - should get no sentinel + let mut buffer = vec![0u8; 1024]; + let n = server.read(&mut buffer).await.unwrap(); + + // Buffer should be empty (no sentinel sent) + let sentinel_json = serde_json::to_vec(&ControlMessage::Sentinel).unwrap(); + assert!( + n == 0 + || !buffer[..n] + .windows(sentinel_json.len()) + .any(|w| w == sentinel_json.as_slice()), + "Buffer should NOT contain sentinel message when context is stopped" + ); + } + + /// Test that handle_writer handles multiple messages correctly + #[tokio::test] + async fn test_handle_writer_multiple_messages() { + let WriterHarness { + server, + framed_writer, + bytes_tx, + bytes_rx, + alive_rx, + controller, + .. + } = writer_harness().await; + + // Send multiple messages + for i in 0..5 { + let test_msg = TwoPartMessage::from_data(Bytes::from(format!("message {}", i))); + bytes_tx.send(test_msg).await.unwrap(); + } + + // Close the sender to trigger normal termination + drop(bytes_tx); + + let result = handle_writer(framed_writer, bytes_rx, alive_rx, controller).await; + + assert!(result.is_ok()); + + // Decode from server side to verify all messages plus sentinel + let mut reader = FramedRead::new(server, TwoPartCodec::default()); + for i in 0..5 { + let msg = recv_msg(&mut reader).await; + assert_data_only_message(msg, format!("message {}", i).as_bytes()); + } + + let sentinel = recv_msg(&mut reader).await; + assert_sentinel_message(sentinel); + } + + /// Test that alive_rx is dropped after handle_writer completes + #[tokio::test] + async fn test_handle_writer_drops_alive_rx() { + let WriterHarness { + framed_writer, + bytes_tx, + bytes_rx, + alive_tx, + alive_rx, + controller, + .. + } = writer_harness().await; + + // Close the sender to trigger normal termination + drop(bytes_tx); + + let result = handle_writer(framed_writer, bytes_rx, alive_rx, controller).await; + + assert!(result.is_ok()); + + // alive_tx should now be closed because alive_rx was dropped + assert!(alive_tx.is_closed()); + } + + /// Test handle_writer with header-only messages (control messages) + #[tokio::test] + async fn test_handle_writer_header_only_messages() { + let WriterHarness { + server, + framed_writer, + bytes_tx, + bytes_rx, + alive_rx, + controller, + .. + } = writer_harness().await; + + // Send a header-only message + let header_msg = TwoPartMessage::from_header(Bytes::from("header content")); + bytes_tx.send(header_msg).await.unwrap(); + + // Close the sender + drop(bytes_tx); + + let result = handle_writer(framed_writer, bytes_rx, alive_rx, controller).await; + + assert!(result.is_ok()); + + let mut reader = FramedRead::new(server, TwoPartCodec::default()); + + let header_msg = recv_msg(&mut reader).await; + assert_header_only_message(header_msg, b"header content"); + + let sentinel = recv_msg(&mut reader).await; + assert_sentinel_message(sentinel); + } + + /// Test handle_writer with mixed header and data messages + #[tokio::test] + async fn test_handle_writer_mixed_messages() { + let WriterHarness { + server, + framed_writer, + bytes_tx, + bytes_rx, + alive_rx, + controller, + .. + } = writer_harness().await; + + // Send mixed messages + bytes_tx + .send(TwoPartMessage::from_header(Bytes::from("header1"))) + .await + .unwrap(); + bytes_tx + .send(TwoPartMessage::from_data(Bytes::from("data1"))) + .await + .unwrap(); + bytes_tx + .send(TwoPartMessage::from_parts( + Bytes::from("header2"), + Bytes::from("data2"), + )) + .await + .unwrap(); + + // Close the sender + drop(bytes_tx); + + let result = handle_writer(framed_writer, bytes_rx, alive_rx, controller).await; + + assert!(result.is_ok()); + + let mut reader = FramedRead::new(server, TwoPartCodec::default()); + + let first = recv_msg(&mut reader).await; + assert_header_only_message(first, b"header1"); + + let second = recv_msg(&mut reader).await; + assert_data_only_message(second, b"data1"); + + let third = recv_msg(&mut reader).await; + assert_header_and_data_message(third, b"header2", b"data2"); + + let sentinel = recv_msg(&mut reader).await; + assert_sentinel_message(sentinel); + } +} diff --git a/lib/runtime/src/pipeline/network/tcp/test_utils.rs b/lib/runtime/src/pipeline/network/tcp/test_utils.rs new file mode 100644 index 000000000..acacca6fe --- /dev/null +++ b/lib/runtime/src/pipeline/network/tcp/test_utils.rs @@ -0,0 +1,21 @@ +// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Test utilities shared across TCP transport tests. + +use tokio::net::TcpListener; + +/// Creates a connected TCP pair for testing. +/// +/// Returns a tuple of (client, server) TcpStream instances that are connected to each other. +/// This is useful for testing functions that operate on TCP streams without needing +/// actual network communication. +pub async fn create_tcp_pair() -> (tokio::net::TcpStream, tokio::net::TcpStream) { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + + let client = tokio::net::TcpStream::connect(addr).await.unwrap(); + let (server, _) = listener.accept().await.unwrap(); + + (client, server) +}