v0.2.8-pre.007-fix.003

crates/ksp-onchain-transport-lib/unit_tests/ws_session.rs

@@ -1,5 +1,5 @@
 // file: crates/ksp-onchain-transport-lib/unit_tests/ws_session.rs
-// version: 10
+// version: 11
 
 use futures_util::SinkExt; // rust-rules: trait-import
 use futures_util::StreamExt; // rust-rules: trait-import
@@ -189,16 +189,17 @@ async fn yield_runtime_steps() {
     return;
 }
 
-async fn wait_for_observed_ping(ping_rx: &mut tokio::sync::mpsc::UnboundedReceiver<()>) {
-    for _ in 0..256 {
-        match ping_rx.try_recv() {
-            std::result::Result::Ok(()) => return,
-            std::result::Result::Err(tokio::sync::mpsc::error::TryRecvError::Empty) => tokio::task::yield_now().await,
-            std::result::Result::Err(tokio::sync::mpsc::error::TryRecvError::Disconnected) => panic!("heartbeat observation channel disconnected before Ping"),
-        }
-    }
-    assert!(matches!(ping_rx.try_recv(), std::result::Result::Ok(())), "heartbeat Ping was not observed after bounded scheduler progress");
-    return;
+async fn wait_for_observed_ping_with_io_progress(ping_rx: &mut tokio::sync::mpsc::UnboundedReceiver<()>) -> std::time::Duration {
+    let started = tokio::time::Instant::now();
+    tokio::time::resume();
+    let observed = tokio::time::timeout(std::time::Duration::from_secs(1), ping_rx.recv()).await;
+    tokio::time::pause();
+    let elapsed = tokio::time::Instant::now().duration_since(started);
+    assert!(
+        matches!(observed, std::result::Result::Ok(std::option::Option::Some(()))),
+        "heartbeat Ping was not observed during bounded real I/O progress"
+    );
+    return elapsed;
 }
 
 #[tokio::test(flavor = "current_thread")]
@@ -469,12 +470,14 @@ async fn helius_heartbeat_sends_ping_at_sixty_seconds_and_rearms() {
     yield_runtime_steps().await;
     assert!(matches!(ping_rx.try_recv(), std::result::Result::Err(tokio::sync::mpsc::error::TryRecvError::Empty)));
     tokio::time::advance(std::time::Duration::from_secs(1)).await;
-    wait_for_observed_ping(&mut ping_rx).await;
-    tokio::time::advance(std::time::Duration::from_secs(59)).await;
+    let first_io_elapsed = wait_for_observed_ping_with_io_progress(&mut ping_rx).await;
+    let pre_second_interval =
+        std::time::Duration::from_secs(59).checked_sub(first_io_elapsed).expect("bounded first Ping I/O progress must remain below one second");
+    tokio::time::advance(pre_second_interval).await;
     yield_runtime_steps().await;
     assert!(matches!(ping_rx.try_recv(), std::result::Result::Err(tokio::sync::mpsc::error::TryRecvError::Empty)));
-    tokio::time::advance(std::time::Duration::from_secs(1)).await;
-    wait_for_observed_ping(&mut ping_rx).await;
+    tokio::time::advance(std::time::Duration::from_secs(1) + first_io_elapsed).await;
+    let _second_io_elapsed = wait_for_observed_ping_with_io_progress(&mut ping_rx).await;
     session.close().await.expect("Helius heartbeat fixture session must close");
     tokio::time::resume();
     server.await.expect("Helius heartbeat fixture server must complete");