Leverage session info API to improve test case for peer timeout

Author: davidsteiner

crates/hotfix/src/session.rs

@@ -32,13 +32,14 @@ use crate::message::sequence_reset::SequenceReset;
 use crate::message::test_request::TestRequest;
 use crate::message_utils::is_admin;
 use crate::session::event::AwaitingActiveSessionResponse;
-use crate::session::info::SessionInfo;
 use crate::session::state::{AwaitingResendState, TestRequestId};
 use crate::session_schedule::SessionSchedule;
 use event::SessionEvent;
 use hotfix_message::parsed_message::ParsedMessage;
 use state::SessionState;
 
+pub use info::{SessionInfo, Status};
+
 const SCHEDULE_CHECK_INTERVAL: u64 = 1;
 
 #[derive(Clone)]

crates/hotfix/src/session/info.rs

@@ -13,7 +13,7 @@ pub struct SessionInfo {
 ///
 /// These roughly correspond to the `SessionState` variants but don't contain
 /// internal state.
-#[derive(Clone, Debug)]
+#[derive(Clone, Debug, PartialEq)]
 pub enum Status {
     AwaitingLogon,
     AwaitingResend,

crates/hotfix/tests/common/mod.rs

@@ -1,3 +1,4 @@
 pub mod mock_application;
 pub mod mock_counterparty;
+pub mod session_assertions;
 pub mod test_messages;

crates/hotfix/tests/common/session_assertions.rs

@@ -0,0 +1,34 @@
+use crate::common::test_messages::TestMessage;
+use hotfix::session::{SessionRef, Status};
+use std::time::Duration;
+
+pub trait SessionAssertions {
+    async fn assert_status(&self, expected_status: Status);
+    async fn assert_status_with_timeout(&self, expected_status: Status, timeout: Duration);
+}
+
+impl SessionAssertions for SessionRef<TestMessage> {
+    async fn assert_status(&self, expected_status: Status) {
+        self.assert_status_with_timeout(expected_status, Duration::from_millis(10))
+            .await;
+    }
+
+    async fn assert_status_with_timeout(&self, expected_status: Status, timeout: Duration) {
+        let deadline = tokio::time::Instant::now() + timeout;
+        let retry_interval = Duration::from_millis(1);
+
+        let mut session_info = self.get_session_info().await;
+        while tokio::time::Instant::now() < deadline {
+            if session_info.status == expected_status {
+                return;
+            }
+            tokio::time::sleep(retry_interval).await;
+            session_info = self.get_session_info().await;
+        }
+
+        let actual_status = session_info.status;
+        panic!(
+            "session did not reach expected status within timeout. Expected: {expected_status:?}, Actual: {actual_status:?}"
+        );
+    }
+}

crates/hotfix/tests/session_tests.rs

@@ -1,9 +1,10 @@
 use crate::common::mock_application::MockApplication;
 use crate::common::mock_counterparty::MockCounterparty;
+use crate::common::session_assertions::SessionAssertions;
 use crate::common::test_messages::TestMessage;
 use hotfix::application::ApplicationRef;
 use hotfix::config::SessionConfig;
-use hotfix::session::SessionRef;
+use hotfix::session::{SessionRef, Status};
 use hotfix::store::in_memory::InMemoryMessageStore;
 use hotfix_message::Part;
 use hotfix_message::fix44::MSG_TYPE;
@@ -56,16 +57,18 @@ async fn test_heartbeats() {
 /// if no response is received within the timeout period.
 #[tokio::test(start_paused = true)]
 async fn test_peer_timeout() {
-    let (_session, mut mock_counterparty) = setup().await;
+    let (session, mut mock_counterparty) = setup().await;
     let peer_interval = (1.2 * HEARTBEAT_INTERVAL as f64) as u64 + 1;
 
     // assert that a logon message is received (type 'A')
     mock_counterparty
         .assert_next(|msg| assert_eq!(msg.header().get::<&str>(MSG_TYPE).unwrap(), "A"))
         .await;
+    session.assert_status(Status::AwaitingLogon).await;
+
     // counterparty responds with a logon to establish a happy session
     mock_counterparty.send_logon().await;
-    tokio::task::yield_now().await;
+    session.assert_status(Status::Active).await;
 
     // let's wait enough time for a heartbeat and assert that the heartbeat was sent
     tokio::time::advance(std::time::Duration::from_secs(HEARTBEAT_INTERVAL + 1)).await;