test: Add first test, and timeout so we dont hang on bad test

Author: Leonidas Loucas

rust/lib/srpc/client/Cargo.toml

@@ -43,4 +43,6 @@ path = "examples/rust_client_example2.rs"
 required-features = []
 
 [dev-dependencies]
-tokio = { version = "1", features = ["full", "macros"] }
+rstest = "0.23.0"
+test-log = { version = "0.2.16", features = ["trace", "color"] }
+tokio = { version = "1", features = ["full", "macros", "test-util"] }

rust/lib/srpc/client/src/lib.rs

@@ -6,7 +6,7 @@ use std::error::Error;
 use std::fmt;
 use std::pin::Pin;
 use std::sync::Arc;
-use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt, BufReader};
+use tokio::io::{AsyncRead, AsyncWrite, AsyncWriteExt, BufReader};
 use tokio::net::TcpStream;
 use tokio::sync::{mpsc, Mutex};
 use tokio::time::{timeout, Duration};
@@ -15,6 +15,8 @@ use tokio_util::codec::{FramedRead, LinesCodec};
 use tracing::debug;
 
 mod chunk_limiter;
+#[cfg(test)]
+mod tests;
 
 // Custom error type
 #[derive(Debug)]
@@ -40,13 +42,19 @@ pub struct ClientConfig {
 pub struct ReceiveOptions {
     channel_buffer_size: usize,
     max_chunk_size: usize,
+    read_next_line_duration: Duration,
 }
 
 impl ReceiveOptions {
-    pub fn new(channel_buffer_size: usize, max_chunk_size: usize) -> Self {
+    pub fn new(
+        channel_buffer_size: usize,
+        max_chunk_size: usize,
+        read_next_line_duration: Duration,
+    ) -> Self {
         ReceiveOptions {
             channel_buffer_size,
             max_chunk_size,
+            read_next_line_duration,
         }
     }
 }
@@ -56,6 +64,7 @@ impl Default for ReceiveOptions {
         ReceiveOptions {
             channel_buffer_size: 100,
             max_chunk_size: 16384,
+            read_next_line_duration: Duration::from_secs(10),
         }
     }
 }
@@ -251,14 +260,15 @@ where
         let stream = Arc::clone(&self.stream);
         let (tx, rx) = mpsc::channel(opts.channel_buffer_size);
         let max_chunk_size = opts.max_chunk_size;
+        let read_next_line_duration = opts.read_next_line_duration;
 
         tokio::spawn(async move {
             let mut guard = stream.lock().await;
             let limited_reader = ChunkLimiter::new(&mut *guard, max_chunk_size);
             let buf_reader = BufReader::new(limited_reader);
             let mut framed = FramedRead::new(buf_reader, LinesCodec::new());
 
-            while let Some(line_res) = framed.next().await {
+            while let Ok(Some(line_res)) = timeout(read_next_line_duration, framed.next()).await {
                 let line_res = line_res.map_err(|e| Box::new(e) as Box<dyn Error + Send>);
 
                 match line_res {

rust/lib/srpc/client/src/tests.rs

@@ -0,0 +1 @@
+mod lib;

rust/lib/srpc/client/src/tests/lib.rs

@@ -0,0 +1,113 @@
+use std::{error::Error, num::NonZeroU8};
+
+use crate::{ClientConfig, ConnectedClient, ReceiveOptions};
+
+use rstest::rstest;
+use tokio::{
+    io::{duplex, AsyncReadExt, AsyncWriteExt, DuplexStream},
+    sync::mpsc,
+};
+
+fn setup_test_client() -> (ConnectedClient<DuplexStream>, DuplexStream) {
+    let (client_stream, server_stream) = duplex(1024);
+
+    let config = ClientConfig::new("example.com", 443, "/", "", "");
+    (ConnectedClient::new(config, client_stream), server_stream)
+}
+
+fn n_message(num: NonZeroU8) -> impl FnMut(&str) -> bool {
+    let mut seen = 0;
+    move |_msg: &str| {
+        if seen + 1 == num.get() {
+            false
+        } else {
+            seen += 1;
+            true
+        }
+    }
+}
+
+fn one_message() -> impl FnMut(&str) -> bool {
+    n_message(NonZeroU8::new(1).unwrap())
+}
+
+async fn check_message(
+    server_message: &str,
+    rx: &mut mpsc::Receiver<Result<String, Box<dyn Error + Send>>>,
+) {
+    if let Some(Ok(received_msg)) = rx.recv().await {
+        assert_eq!(received_msg, server_message.trim());
+    } else {
+        panic!("Did not receive expected message from server");
+    }
+}
+
+async fn check_server(
+    client_message: &str,
+    server_stream: &mut DuplexStream,
+) -> Result<(), Box<dyn Error>> {
+    let mut server_buf = vec![0u8; client_message.len()];
+    server_stream.read_exact(&mut server_buf).await?;
+    assert_eq!(&server_buf, client_message.as_bytes());
+    Ok(())
+}
+
+#[test_log::test(rstest)]
+#[tokio::test(start_paused = true)]
+async fn test_connected_client_send_and_receive() -> Result<(), Box<dyn Error>> {
+    let (connected_client, mut server_stream) = setup_test_client();
+
+    let client_message = "Hello from client\n";
+    connected_client.send_message(client_message).await?;
+
+    check_server(client_message, &mut server_stream).await?;
+
+    let server_message = "Hello from server\n";
+    server_stream.write_all(server_message.as_bytes()).await?;
+
+    let should_continue = one_message();
+
+    let opts = ReceiveOptions::default();
+    let mut rx = connected_client
+        .receive_message(false, should_continue, &opts)
+        .await?;
+
+    check_message(server_message, &mut rx).await;
+
+    Ok(())
+}
+
+#[test_log::test(rstest)]
+#[tokio::test(start_paused = true)]
+async fn test_connected_client_send_and_receive_stream() -> Result<(), Box<dyn Error>> {
+    let (connected_client, mut server_stream) = setup_test_client();
+
+    let client_message = "Hello from client\n";
+    connected_client.send_message(client_message).await?;
+
+    check_server(client_message, &mut server_stream).await?;
+
+    server_stream.write_all("\n".as_bytes()).await?;
+
+    let should_continue = one_message();
+
+    let opts = ReceiveOptions::default();
+    let mut rx = connected_client
+        .receive_message(true, should_continue, &opts)
+        .await?;
+
+    check_message("", &mut rx).await;
+
+    server_stream.write_all("first\n".as_bytes()).await?;
+
+    server_stream.write_all("second\n".as_bytes()).await?;
+
+    let should_continue = n_message(NonZeroU8::new(2).unwrap());
+    let mut rx = connected_client
+        .receive_message(false, should_continue, &opts)
+        .await?;
+
+    check_message("first", &mut rx).await;
+    check_message("second", &mut rx).await;
+    Ok(())
+}