From 9001b440793616782839af893c3eb185f805d1f2 Mon Sep 17 00:00:00 2001
From: Youngjoon Lee <5462944+youngjoon-lee@users.noreply.github.com>
Date: Fri, 20 Dec 2024 20:36:34 +0900
Subject: [PATCH] test(blend): add MockBlendMessage::payload function to get
 payload without unpacking (#967)

---
 nomos-blend/message/src/mock/mod.rs | 29 ++++++++++++++++++-----------
 1 file changed, 18 insertions(+), 11 deletions(-)

diff --git a/nomos-blend/message/src/mock/mod.rs b/nomos-blend/message/src/mock/mod.rs
index ed49a516..4b2b12a1 100644
--- a/nomos-blend/message/src/mock/mod.rs
+++ b/nomos-blend/message/src/mock/mod.rs
@@ -75,17 +75,7 @@ impl BlendMessage for MockBlendMessage {
 
         // If this is the last layer
         if message[NODE_ID_SIZE..NODE_ID_SIZE * 2] == DUMMY_NODE_ID {
-            let padded_payload = &message[NODE_ID_SIZE * MAX_LAYERS..];
-            // remove the payload padding
-            match padded_payload
-                .iter()
-                .rposition(|&x| x == PAYLOAD_PADDING_SEPARATOR)
-            {
-                Some(pos) => {
-                    return Ok((padded_payload[0..pos].to_vec(), true));
-                }
-                _ => return Err(Error::InvalidBlendMessage),
-            }
+            return Ok((Self::payload(message)?, true));
         }
 
         let mut new_message: Vec<u8> = Vec::with_capacity(MESSAGE_SIZE);
@@ -96,6 +86,20 @@ impl BlendMessage for MockBlendMessage {
     }
 }
 
+impl MockBlendMessage {
+    pub fn payload(message: &[u8]) -> Result<Vec<u8>, Error> {
+        let padded_payload = &message[NODE_ID_SIZE * MAX_LAYERS..];
+        // remove the payload padding
+        match padded_payload
+            .iter()
+            .rposition(|&x| x == PAYLOAD_PADDING_SEPARATOR)
+        {
+            Some(pos) => Ok(padded_payload[0..pos].to_vec()),
+            _ => Err(Error::InvalidBlendMessage),
+        }
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -106,16 +110,19 @@ mod tests {
         let payload = [7; 10];
         let message = MockBlendMessage::build_message(&payload, &node_ids).unwrap();
         assert_eq!(message.len(), MESSAGE_SIZE);
+        assert_eq!(MockBlendMessage::payload(&message).unwrap(), payload);
 
         let (message, is_fully_unwrapped) =
             MockBlendMessage::unwrap_message(&message, &node_ids[0]).unwrap();
         assert!(!is_fully_unwrapped);
         assert_eq!(message.len(), MESSAGE_SIZE);
+        assert_eq!(MockBlendMessage::payload(&message).unwrap(), payload);
 
         let (message, is_fully_unwrapped) =
             MockBlendMessage::unwrap_message(&message, &node_ids[1]).unwrap();
         assert!(!is_fully_unwrapped);
         assert_eq!(message.len(), MESSAGE_SIZE);
+        assert_eq!(MockBlendMessage::payload(&message).unwrap(), payload);
 
         let (unwrapped_payload, is_fully_unwrapped) =
             MockBlendMessage::unwrap_message(&message, &node_ids[2]).unwrap();