Parse messages in session tests so assertions can be made on their details

Author: davidsteiner

crates/hotfix/tests/common/mock_counterparty.rs

@@ -1,25 +1,30 @@
-use hotfix::message::FixMessage;
+use hotfix::message::{FixMessage, RawFixMessage};
 use hotfix::session::SessionRef;
 use hotfix::transport::FixConnection;
 use hotfix::transport::reader::ReaderRef;
 use hotfix::transport::writer::{WriterMessage, WriterRef};
+use hotfix_message::dict::Dictionary;
+use hotfix_message::message::{Config as MessageConfig, Message};
+use hotfix_message::parsed_message::ParsedMessage;
 use std::time::Duration;
+use tokio::sync::mpsc::Receiver;
 use tokio::sync::{mpsc, oneshot};
 
-const DEFAULT_TIMEOUT: Duration = Duration::from_secs(1);
+const DEFAULT_TIMEOUT: Duration = Duration::from_millis(10);
 
 pub struct MockCounterparty {
-    // Receiver-End of the channel
-    receiver: mpsc::UnboundedReceiver<WriterMessage>,
-    // History of Received messages from the client
-    messages: Vec<WriterMessage>,
+    receiver: Receiver<WriterMessage>,
+    // History of received messages from the session
+    messages: Vec<Message>,
+    dictionary: Dictionary,
+    message_config: MessageConfig,
     _connection: FixConnection,
     _dc_sender: oneshot::Sender<()>,
 }
 
 impl MockCounterparty {
     pub async fn start(session_ref: SessionRef<impl FixMessage>) -> Self {
-        let (writer_ref, receiver) = Self::spawn_writer();
+        let (writer_ref, receiver) = Self::create_writer();
         let (reader_ref, dc_sender) = Self::create_reader();
         let connection = FixConnection::new(writer_ref, reader_ref);
 
@@ -28,65 +33,92 @@ impl MockCounterparty {
         Self {
             receiver,
             messages: vec![],
+            dictionary: Dictionary::fix44(),
+            message_config: MessageConfig::default(),
             _connection: connection,
             _dc_sender: dc_sender,
         }
     }
 
-    /// Listen to the next message on the channel
-    pub async fn get_next(&mut self, timeout: Option<Duration>) -> &WriterMessage {
-        let timeout = timeout.unwrap_or(DEFAULT_TIMEOUT);
-        let msg = tokio::time::timeout(timeout, self.receiver.recv())
+    /// Waits for and returns the next message received from the session.
+    ///
+    /// A `None` response indicates we have been disconnected, either through the channel
+    /// dropping on the session's side, or through an explicit `Disconnect` message.
+    async fn get_next(&mut self) -> Option<&Message> {
+        self.receiver
+            .recv()
             .await
-            .unwrap_or_else(|_| panic!("Message not received in less than {timeout:?}"))
-            .unwrap_or_else(|| panic!("Received message is None"));
-        self.messages.push(msg);
-        self.messages.last().unwrap()
+            .and_then(|writer_message| match writer_message {
+                WriterMessage::SendMessage(raw_message) => {
+                    let message = self.parse_message(&raw_message);
+                    self.messages.push(message);
+                    self.messages.last()
+                }
+                WriterMessage::Disconnect => None,
+            })
     }
 
-    pub async fn assert_next<F>(&mut self, timeout: Option<Duration>, assertion: F)
+    fn parse_message(&self, raw_message: &RawFixMessage) -> Message {
+        match Message::from_bytes(
+            &self.message_config,
+            &self.dictionary,
+            raw_message.as_bytes(),
+        ) {
+            ParsedMessage::Valid(valid_message) => valid_message,
+            _ => {
+                panic!("only valid messages are supported in the mock counterparty")
+            }
+        }
+    }
+
+    pub async fn assert_next<F>(&mut self, assertion: F)
     where
-        F: FnOnce(&WriterMessage) -> bool,
+        F: FnOnce(&Message) -> bool,
     {
-        let msg = self.get_next(timeout).await;
-        assert!(assertion(msg));
+        self.assert_next_with_timeout(assertion, DEFAULT_TIMEOUT)
+            .await;
     }
 
-    pub async fn assert_disconnected(&mut self, timeout: Option<Duration>) {
-        self.assert_next(timeout, |msg| matches!(msg, &WriterMessage::Disconnect))
-            .await;
-        assert!(self.receiver.is_closed());
+    pub async fn assert_next_with_timeout<F>(&mut self, assertion: F, timeout: Duration)
+    where
+        F: FnOnce(&Message) -> bool,
+    {
+        match tokio::time::timeout(timeout, self.get_next()).await {
+            Ok(Some(message)) => {
+                assertion(message);
+            }
+            Ok(None) => {
+                panic!("disconnected before receiving any message");
+            }
+            Err(_) => {
+                panic!("timeout expired before receiving any message");
+            }
+        }
+    }
+
+    pub async fn assert_disconnected(&mut self) {
+        self.assert_disconnected_with_timeout(DEFAULT_TIMEOUT).await;
     }
 
-    fn spawn_writer() -> (WriterRef, mpsc::UnboundedReceiver<WriterMessage>) {
-        // mock implementation to receive messages from the session
-        // and hold on to them for test assertions
-        let (tx, rx) = mpsc::unbounded_channel();
-        let (sender, mailbox) = mpsc::channel(10);
-        tokio::spawn(Self::run_writer(mailbox, tx));
+    pub async fn assert_disconnected_with_timeout(&mut self, timeout: Duration) {
+        if tokio::time::timeout(timeout, async {
+            // keep consuming messages until a disconnect occurs
+            while self.get_next().await.is_some() {}
+        })
+        .await
+        .is_err()
+        {
+            panic!("timeout expired before a disconnect occurred");
+        }
+    }
 
-        (WriterRef::new(sender), rx)
+    fn create_writer() -> (WriterRef, Receiver<WriterMessage>) {
+        let (sender, receiver) = mpsc::channel(10);
+        (WriterRef::new(sender), receiver)
     }
 
     fn create_reader() -> (ReaderRef, oneshot::Sender<()>) {
         let (dc_sender, dc_receiver) = oneshot::channel();
         (ReaderRef::new(dc_receiver), dc_sender)
     }
-
-    async fn run_writer(
-        mut mailbox: mpsc::Receiver<WriterMessage>,
-        received_messages: mpsc::UnboundedSender<WriterMessage>,
-    ) {
-        while let Some(msg) = mailbox.recv().await {
-            println!("Received message from session: {msg:?}");
-            let disconnect = matches!(msg, WriterMessage::Disconnect);
-            if let Err(e) = received_messages.send(msg) {
-                panic!("Failed to send message. Error: {e:?}");
-            }
-            if disconnect {
-                println!("Disconnecting");
-                break;
-            }
-        }
-    }
 }

crates/hotfix/tests/session_tests.rs

@@ -5,18 +5,25 @@ use hotfix::application::ApplicationRef;
 use hotfix::config::SessionConfig;
 use hotfix::session::SessionRef;
 use hotfix::store::in_memory::InMemoryMessageStore;
+use hotfix_message::Part;
+use hotfix_message::fix44::MSG_TYPE;
 
 mod common;
 
 #[tokio::test]
 async fn test_happy_login_flow() {
     let session = create_session();
     let mut mock_counterparty = MockCounterparty::start(session.clone()).await;
-    mock_counterparty.assert_next(None, |_msg| true).await;
+
+    // assert that a logon message is received (type 'A')
+    mock_counterparty
+        .assert_next(|msg| msg.header().get::<&str>(MSG_TYPE).unwrap() == "A")
+        .await;
+
     session
         .disconnect("Test Session Finished".to_string())
         .await;
-    mock_counterparty.assert_disconnected(None).await;
+    mock_counterparty.assert_disconnected().await;
 }
 
 fn create_session() -> SessionRef<TestMessage> {