From 0a9054a51d0718cf840ba9f84f14c351991e046c Mon Sep 17 00:00:00 2001
From: Nils Andreas Svee <me@lochnair.net>
Date: Wed, 29 Apr 2026 01:14:32 +0700
Subject: [PATCH] test(managed): improve test reliability and coverage

---
 crates/irtt-client/src/managed/hub.rs    | 20 +++++++++++++++
 crates/irtt-client/src/managed/runner.rs | 31 ++++++++++++++++++------
 2 files changed, 44 insertions(+), 7 deletions(-)

diff --git a/crates/irtt-client/src/managed/hub.rs b/crates/irtt-client/src/managed/hub.rs
index e3aabfa..1aec013 100644
--- a/crates/irtt-client/src/managed/hub.rs
+++ b/crates/irtt-client/src/managed/hub.rs
@@ -148,6 +148,15 @@ impl EventHub {
             subscriber.disconnect();
         }
     }
+
+    #[cfg(test)]
+    fn subscriber_count(&self) -> usize {
+        self.inner
+            .subscribers
+            .lock()
+            .expect("event hub mutex poisoned")
+            .len()
+    }
 }
 
 impl Default for EventHub {
@@ -363,6 +372,17 @@ mod tests {
         assert!(sub.try_recv().unwrap().is_none());
     }
 
+    #[test]
+    fn dropping_subscription_unregisters_from_hub() {
+        let hub = EventHub::new();
+        let sub = hub.subscribe(SubscriberConfig::default()).unwrap();
+        assert_eq!(hub.subscriber_count(), 1);
+
+        drop(sub);
+
+        assert_eq!(hub.subscriber_count(), 0);
+    }
+
     #[test]
     fn disconnect_all_wakes_blocking_receivers() {
         let hub = EventHub::new();
diff --git a/crates/irtt-client/src/managed/runner.rs b/crates/irtt-client/src/managed/runner.rs
index 96ecdd9..d50d6af 100644
--- a/crates/irtt-client/src/managed/runner.rs
+++ b/crates/irtt-client/src/managed/runner.rs
@@ -329,8 +329,8 @@ mod tests {
                 let ts = TimestampFields {
                     recv_wall: Some(1_000_000_000),
                     recv_mono: Some(100_000),
-                    send_wall: Some(1_000_100_000),
-                    send_mono: Some(200_000),
+                    send_wall: Some(1_000_000_000),
+                    send_mono: Some(100_000),
                     ..Default::default()
                 };
                 socket
@@ -408,20 +408,37 @@ mod tests {
         packet
     }
 
+    fn recv_event_with_timeout(sub: &EventSubscription) -> ClientEvent {
+        let deadline = Instant::now() + Duration::from_secs(2);
+        loop {
+            match sub.try_recv() {
+                Ok(Some(event)) => return event,
+                Ok(None) if Instant::now() < deadline => {
+                    thread::sleep(Duration::from_millis(1));
+                }
+                Ok(None) => panic!("timed out waiting for managed client event"),
+                Err(err) => panic!("subscription ended while waiting for event: {err}"),
+            }
+        }
+    }
+
     fn collect_until_closed(sub: &EventSubscription) -> Vec<ClientEvent> {
         let mut events = Vec::new();
-        loop {
-            match sub.recv() {
-                Ok(event) => {
+        let deadline = Instant::now() + Duration::from_secs(2);
+        while Instant::now() < deadline {
+            match sub.try_recv() {
+                Ok(Some(event)) => {
                     let closed = matches!(event, ClientEvent::SessionClosed { .. });
                     events.push(event);
                     if closed {
                         return events;
                     }
                 }
+                Ok(None) => thread::sleep(Duration::from_millis(1)),
                 Err(_) => return events,
             }
         }
+        panic!("timed out waiting for session close event");
     }
 
     #[test]
@@ -477,7 +494,7 @@ mod tests {
             .unwrap();
 
         loop {
-            if matches!(sub.recv().unwrap(), ClientEvent::EchoReply { .. }) {
+            if matches!(recv_event_with_timeout(&sub), ClientEvent::EchoReply { .. }) {
                 break;
             }
         }
@@ -530,7 +547,7 @@ mod tests {
         let mut cfg = config(addr, Some(Duration::from_millis(10)));
         cfg.run_mode = RunMode::NoTest;
         let session = ManagedClient::start(cfg).unwrap();
-        assert!(rx.recv().unwrap());
+        assert!(rx.recv_timeout(Duration::from_secs(2)).unwrap());
         let outcome = session.join().unwrap();
         done.join().unwrap();