From b4d56e7d59ee66451830f334039e8d0724bc26f3 Mon Sep 17 00:00:00 2001
From: fys <fengys1996@gmail.com>
Date: Tue, 24 Dec 2024 14:03:31 +0800
Subject: [PATCH] fix: add tests and fix some

1. add test_ringbuf_spsc_with_wait_result_and_notify
2. calc data_len of datablock may panic
---
 src/ringbuf/data_block.rs | 10 ++++++----
 tests/common.rs           |  5 +++--
 tests/ringbuf_spsc.rs     | 18 ++++++++++++++----
 3 files changed, 23 insertions(+), 10 deletions(-)

diff --git a/src/ringbuf/data_block.rs b/src/ringbuf/data_block.rs
index a381abd..c1518cd 100644
--- a/src/ringbuf/data_block.rs
+++ b/src/ringbuf/data_block.rs
@@ -100,17 +100,19 @@ impl<T> DataBlock<T> {
         object: Arc<T>,
     ) -> Result<Self> {
         let header_len_u32 = convert_num!(HEADER_LEN, u32)?;
-        let data_len = len - header_len_u32;
 
         ensure!(
-            data_len > 0,
+            len >= header_len_u32,
             error::InvalidParameterSnafu {
-                detail: "Total_length must be greater than HEADER_LEN.",
+                detail: "Total length must be greater than HEADER_LEN.",
             }
         );
 
+        // Unwrap safety: checked above.
+        let data_len = len.checked_sub(header_len_u32).unwrap();
+
         let header = Header::fow_raw(start_ptr);
-        header.set_capacity(len - header_len_u32);
+        header.set_capacity(data_len);
         header.set_written(0);
         header.set_busy(true);
         header.set_req_id(req_id);
diff --git a/tests/common.rs b/tests/common.rs
index cb4f416..9f0f619 100644
--- a/tests/common.rs
+++ b/tests/common.rs
@@ -103,9 +103,10 @@ pub async fn wait_consumer_online(
     Err("wait consumer online timeout".to_string())
 }
 
-pub fn gen_str(max_len: usize) -> String {
+pub fn gen_str(min_len: usize, max_len: usize) -> String {
     const CHARSET: &[u8] = b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789)(*&^%$#@!~";
-    let len = rand::random::<usize>() % max_len + 1;
+
+    let len = rand::random::<usize>() % (max_len - min_len) + min_len;
 
     let mut s = String::new();
 
diff --git a/tests/ringbuf_spsc.rs b/tests/ringbuf_spsc.rs
index 6bc8d5b..dd315fb 100644
--- a/tests/ringbuf_spsc.rs
+++ b/tests/ringbuf_spsc.rs
@@ -13,24 +13,34 @@ use shm_ringbuf::producer::RingbufProducer;
 
 #[tokio::test]
 async fn test_ringbuf_spsc_base() {
-    do_test_ringbuf_spsc(false, Duration::from_millis(10), None).await;
+    do_test_ringbuf_spsc(false, Duration::from_millis(10), None, 0, 1000).await;
 }
 
 #[tokio::test]
 async fn test_ringbuf_spsc_with_notify() {
     // Set too long interval for testing notify.
-    do_test_ringbuf_spsc(false, Duration::from_secs(100), Some(1000)).await;
+    do_test_ringbuf_spsc(false, Duration::from_secs(100), Some(500), 499, 1000)
+        .await;
 }
 
 #[tokio::test]
 async fn test_ringbuf_spsc_with_wait_result() {
-    do_test_ringbuf_spsc(true, Duration::from_millis(10), None).await;
+    do_test_ringbuf_spsc(true, Duration::from_millis(10), None, 0, 1000).await;
+}
+
+#[tokio::test]
+async fn test_ringbuf_spsc_with_wait_result_and_notify() {
+    // Set too long interval for testing notify.
+    do_test_ringbuf_spsc(true, Duration::from_secs(100), Some(500), 499, 1000)
+        .await;
 }
 
 async fn do_test_ringbuf_spsc(
     wait_result: bool,
     process_interval: Duration,
     notify_limit: Option<u32>,
+    min_msg_len: usize,
+    max_msg_len: usize,
 ) {
     tracing_subscriber::fmt::init();
 
@@ -74,7 +84,7 @@ async fn do_test_ringbuf_spsc(
         };
 
         for i in 0..msg_num {
-            let write_str = gen_str(1000);
+            let write_str = gen_str(min_msg_len, max_msg_len);
 
             expected_send.send(write_str.clone()).await.unwrap();