From 9cd6b9c3e390eb9a91712884377e8c342708918d Mon Sep 17 00:00:00 2001 From: mhossamr Date: Mon, 20 Jul 2026 13:44:22 +0300 Subject: [PATCH] feat(tests): add unit tests for FEO Scheduler --- src/feo/BUILD.bazel | 7 +- src/feo/src/scheduler.rs | 143 +++++++++++++++++++++++++++++++++++++++ 2 files changed, 149 insertions(+), 1 deletion(-) diff --git a/src/feo/BUILD.bazel b/src/feo/BUILD.bazel index efd8a9f..651e708 100644 --- a/src/feo/BUILD.bazel +++ b/src/feo/BUILD.bazel @@ -11,7 +11,7 @@ # SPDX-License-Identifier: Apache-2.0 # ******************************************************************************* -load("@rules_rust//rust:defs.bzl", "rust_library") +load("@rules_rust//rust:defs.bzl", "rust_library", "rust_test") rust_library( name = "libfeo_rust", @@ -103,3 +103,8 @@ cc_library( ], alwayslink = True, ) + +rust_test( + name = "feo_test", + crate = ":libfeo_rust", +) diff --git a/src/feo/src/scheduler.rs b/src/feo/src/scheduler.rs index 12d7639..136f864 100644 --- a/src/feo/src/scheduler.rs +++ b/src/feo/src/scheduler.rs @@ -435,3 +435,146 @@ mod loop_duration_meter { } } } + +#[cfg(test)] +mod tests { + use super::*; + use alloc::vec; // Import the vec! macro + use std::collections::VecDeque; + use std::sync::Mutex; + + #[derive(Default)] + struct MockState { + sent_signals: Vec<(ActivityId, Signal)>, + broadcasted_signals: Vec, + receive_queue: VecDeque, + connected_agents: Vec, + } + + struct MockConnector { + state: Arc>, + } + + impl ConnectScheduler for MockConnector { + fn sync_time(&mut self) -> Result<(), Error> { + Ok(()) + } + + fn send_to_activity(&mut self, id: ActivityId, signal: &Signal) -> Result<(), Error> { + self.state.lock().unwrap().sent_signals.push((id, signal.clone())); + Ok(()) + } + fn connect_remotes(&mut self) -> Result<(), Error> { + Ok(()) + } + + fn broadcast_terminate(&mut self, signal: &Signal) -> Result<(), Error> { + self.state.lock().unwrap().broadcasted_signals.push(signal.clone()); + Ok(()) + } + + fn receive(&mut self, _timeout: feo_time::Duration) -> Result, Error> { + Ok(self.state.lock().unwrap().receive_queue.pop_front()) + } + + fn get_connected_agent_ids(&self) -> Vec { + self.state.lock().unwrap().connected_agents.clone() + } + } + + static INIT: std::sync::Once = std::sync::Once::new(); + + fn init_test_env() { + INIT.call_once(|| { + crate::timestamp::initialize(); + }); + } + + #[test] + fn test_step_ready_activities() { + init_test_env(); + let state = Arc::new(Mutex::new(MockState::default())); + let connector = Box::new(MockConnector { state: state.clone() }); + let shutdown_req = Arc::new(AtomicBool::new(false)); + + let id_a = ActivityId::from(1); + let id_b = ActivityId::from(2); + let id_c = ActivityId::from(3); + + let mut depends = HashMap::new(); + depends.insert(id_a, vec![]); + depends.insert(id_b, vec![]); + depends.insert(id_c, vec![id_a, id_b]); + + let mut scheduler = Scheduler::new( + AgentId::from(1), + feo_time::Duration::from_millis(10), + feo_time::Duration::from_millis(10), + feo_time::Duration::from_millis(10), + depends, + connector, + shutdown_req, + ); + + // Mark A and B as ready (their fake operations have completed) + scheduler.activity_states.get_mut(&id_a).unwrap().triggered = true; + scheduler.activity_states.get_mut(&id_a).unwrap().ready = true; + scheduler.activity_states.get_mut(&id_b).unwrap().triggered = true; + scheduler.activity_states.get_mut(&id_b).unwrap().ready = true; + + // Trigger evaluation + scheduler.step_ready_activities(); + + // C should now be triggered since A and B are ready + assert!(scheduler.activity_states.get(&id_c).unwrap().triggered); + + // Verify the scheduler actually attempted to send the `Step` signal to C over the connector + let sent = &state.lock().unwrap().sent_signals; + assert_eq!(sent.len(), 1); + assert_eq!(sent[0].0, id_c); + assert!(matches!(sent[0].1, Signal::Step(_))); + } + + #[test] + fn test_graceful_shutdown_only_started_activities() { + init_test_env(); + let state = Arc::new(Mutex::new(MockState::default())); + let connector = Box::new(MockConnector { state: state.clone() }); + let shutdown_req = Arc::new(AtomicBool::new(false)); + + let id_1 = ActivityId::from(1); + let id_2 = ActivityId::from(2); + let id_3 = ActivityId::from(3); + + let mut depends = HashMap::new(); + depends.insert(id_1, vec![]); + depends.insert(id_2, vec![]); + depends.insert(id_3, vec![]); + + let mut scheduler = Scheduler::new( + AgentId::from(1), + feo_time::Duration::from_millis(10), + feo_time::Duration::from_millis(10), + feo_time::Duration::from_millis(10), + depends, + connector, + shutdown_req, + ); + + // Emulate a state where activities 1 and 2 successfully started, but 3 never did + scheduler.activity_states.get_mut(&id_1).unwrap().ever_ready = true; + scheduler.activity_states.get_mut(&id_2).unwrap().ever_ready = true; + scheduler.activity_states.get_mut(&id_3).unwrap().ever_ready = false; + + scheduler.shutdown_gracefully("test shutdown"); + + let sent = &state.lock().unwrap().sent_signals; + + // Should only send shutdown to 1 and 2, but NOT 3 + assert_eq!(sent.len(), 2); + let sent_ids: BTreeSet<_> = sent.iter().map(|(id, _)| *id).collect(); + assert!(sent_ids.contains(&id_1)); + assert!(sent_ids.contains(&id_2)); + assert!(!sent_ids.contains(&id_3)); + } +}