diff --git a/app/src/ai/agent_sdk/driver_tests.rs b/app/src/ai/agent_sdk/driver_tests.rs index ae36cf4a..965ff590 100644 --- a/app/src/ai/agent_sdk/driver_tests.rs +++ b/app/src/ai/agent_sdk/driver_tests.rs @@ -1,6 +1,7 @@ use std::{collections::HashMap, ffi::OsString, sync::Arc, time::Duration}; use futures::channel::oneshot; +use instant::Instant; use warp_cli::agent::Harness; use warp_cli::{ OZ_CLI_ENV, OZ_HARNESS_ENV, OZ_PARENT_RUN_ID_ENV, OZ_RUN_ID_ENV, SERVER_ROOT_URL_OVERRIDE_ENV, @@ -118,6 +119,22 @@ fn test_normalize_sse_server_with_headers() { // ── IdleTimeoutSender tests ────────────────────────────────────────────────────── +fn recv_within(rx: &mut oneshot::Receiver, timeout: Duration) -> Option { + let started_at = Instant::now(); + + loop { + if let Some(value) = rx.try_recv().unwrap() { + return Some(value); + } + + if started_at.elapsed() >= timeout { + return None; + } + + std::thread::sleep(Duration::from_millis(10)); + } +} + #[test] fn idle_timeout_sender_send_now_delivers_value() { let (tx, mut rx) = oneshot::channel::(); @@ -144,8 +161,7 @@ fn idle_timeout_sender_send_after_delivers_after_timeout() { // Not yet delivered. assert_eq!(rx.try_recv().unwrap(), None); - std::thread::sleep(Duration::from_millis(100)); - assert_eq!(rx.try_recv().unwrap(), Some(99)); + assert_eq!(recv_within(&mut rx, Duration::from_secs(1)), Some(99)); } #[test] @@ -180,8 +196,7 @@ fn idle_timeout_sender_later_send_after_supersedes_earlier() { // Second timer: short timeout. The first is implicitly cancelled. idle_timeout.end_run_after(Duration::from_millis(50), 2); - std::thread::sleep(Duration::from_millis(100)); - assert_eq!(rx.try_recv().unwrap(), Some(2)); + assert_eq!(recv_within(&mut rx, Duration::from_secs(1)), Some(2)); } #[test]