From 15b9157e9c0ad6f4d3aa54fa874f950e74d6e46b Mon Sep 17 00:00:00 2001
From: Warittorn Cheevachaipimol <warittorn@bandprotocol.com>
Date: Fri, 22 Mar 2024 15:22:56 +0700
Subject: [PATCH] change forwarding to receive message instead of string

---
 examples/forwarding_messages.rs |  7 ++++---
 src/ws_mock_server.rs           | 24 +++++++++++++-----------
 2 files changed, 17 insertions(+), 14 deletions(-)

diff --git a/examples/forwarding_messages.rs b/examples/forwarding_messages.rs
index af9b02a..212b401 100644
--- a/examples/forwarding_messages.rs
+++ b/examples/forwarding_messages.rs
@@ -2,6 +2,7 @@ use futures_util::StreamExt;
 use std::time::Duration;
 use tokio::sync::mpsc;
 use tokio_tungstenite::connect_async;
+use tokio_tungstenite::tungstenite::Message;
 use ws_mock::utils::collect_all_messages;
 use ws_mock::ws_mock_server::{WsMock, WsMockServer};
 
@@ -9,7 +10,7 @@ use ws_mock::ws_mock_server::{WsMock, WsMockServer};
 pub async fn main() {
     let server = WsMockServer::start().await;
 
-    let (mpsc_send, mpsc_recv) = mpsc::channel::<String>(32);
+    let (mpsc_send, mpsc_recv) = mpsc::channel::<Message>(32);
 
     WsMock::new()
         .forward_from_channel(mpsc_recv)
@@ -22,8 +23,8 @@ pub async fn main() {
 
     let (_send, ws_recv) = stream.split();
 
-    mpsc_send.send("message-1".to_string()).await.unwrap();
-    mpsc_send.send("message-2".into()).await.unwrap();
+    mpsc_send.send(Message::Text("message-1".to_string())).await.unwrap();
+    mpsc_send.send(Message::Text("message-2".to_string())).await.unwrap();
 
     let received = collect_all_messages(ws_recv, Duration::from_millis(250)).await;
 
diff --git a/src/ws_mock_server.rs b/src/ws_mock_server.rs
index 3f86ade..6a3d879 100644
--- a/src/ws_mock_server.rs
+++ b/src/ws_mock_server.rs
@@ -56,7 +56,7 @@ const INCOMPLETE_MOCK_PANIC: &str = "A mock must have a response or expected num
 pub struct WsMock {
     matchers: Vec<Box<dyn Matcher>>,
     response_data: Option<String>,
-    forwarding_channel: Option<MpscReceiver<String>>,
+    forwarding_channel: Option<MpscReceiver<Message>>,
     expected_calls: Option<usize>,
     calls: usize,
 }
@@ -123,9 +123,10 @@ impl WsMock {
     ///     use std::time::Duration;
     ///     use futures_util::SinkExt;
     ///     use tokio_tungstenite::connect_async;
+    ///     use tokio_tungstenite::tungstenite::Message;
     ///     let server = WsMockServer::start().await;
     ///
-    ///     let (mpsc_send, mpsc_recv) = mpsc::channel::<String>(32);
+    ///     let (mpsc_send, mpsc_recv) = mpsc::channel::<Message>(32);
     ///
     ///     WsMock::new()
     ///         .forward_from_channel(mpsc_recv)
@@ -138,8 +139,8 @@ impl WsMock {
     ///
     ///     let (_send, ws_recv) = stream.split();
     ///
-    ///     mpsc_send.send("message-1".to_string()).await.unwrap();
-    ///     mpsc_send.send("message-2".into()).await.unwrap();
+    ///     mpsc_send.send(Message::Text("message-1".to_string())).await.unwrap();
+    ///     mpsc_send.send(Message::Text("message-2".into())).await.unwrap();
     ///
     ///     let received = collect_all_messages(ws_recv, Duration::from_millis(250)).await;
     ///
@@ -150,7 +151,7 @@ impl WsMock {
     ///
     /// [`tokio::sync:mpsc`]: https://docs.rs/tokio/1.36.0/tokio/sync/mpsc/index.html
     /// [`Receiver`]: https://docs.rs/tokio/1.36.0/tokio/sync/mpsc/struct.Receiver.html
-    pub fn forward_from_channel(mut self, receiver: MpscReceiver<String>) -> Self {
+    pub fn forward_from_channel(mut self, receiver: MpscReceiver<Message>) -> Self {
         self.forwarding_channel = Some(receiver);
         self
     }
@@ -398,11 +399,11 @@ impl WsMockServer {
     /// sender.
     #[doc(hidden)]
     async fn forward_messages_task(
-        mut incoming: MpscReceiver<String>,
+        mut incoming: MpscReceiver<Message>,
         outgoing: MpscSender<Message>,
     ) {
         while let Some(msg) = incoming.recv().await {
-            outgoing.send(Message::text(msg)).await.unwrap();
+            outgoing.send(msg).await.unwrap();
         }
     }
 
@@ -487,6 +488,7 @@ mod tests {
     use std::time::Duration;
     use tokio::sync::mpsc;
     use tokio_tungstenite::connect_async;
+    use tokio_tungstenite::tungstenite::Message;
 
     #[tokio::test]
     async fn test_wss_mockserver() {
@@ -519,7 +521,7 @@ mod tests {
     async fn test_forwarding_channel() {
         let server = WsMockServer::start().await;
 
-        let (mpsc_send, mpsc_recv) = mpsc::channel::<String>(32);
+        let (mpsc_send, mpsc_recv) = mpsc::channel::<Message>(32);
 
         WsMock::new()
             .matcher(Any::new())
@@ -533,8 +535,8 @@ mod tests {
 
         let (_send, ws_recv) = stream.split();
 
-        mpsc_send.send("message-1".to_string()).await.unwrap();
-        mpsc_send.send("message-2".into()).await.unwrap();
+        mpsc_send.send(Message::Text("message-1".to_string())).await.unwrap();
+        mpsc_send.send(Message::Text("message-2".into())).await.unwrap();
 
         let received = collect_all_messages(ws_recv, Duration::from_millis(250)).await;
 
@@ -546,7 +548,7 @@ mod tests {
     async fn test_shutdown_with_active_channel() {
         let mut server = WsMockServer::start().await;
 
-        let (_, mpsc_recv) = mpsc::channel::<String>(32);
+        let (_, mpsc_recv) = mpsc::channel::<Message>(32);
 
         WsMock::new()
             .matcher(Any::new())