From f89086d8e40c377921576e04038d7a08680f7907 Mon Sep 17 00:00:00 2001
From: Martin Pitt <martin@piware.de>
Date: Fri, 9 Dec 2022 09:59:39 +0100
Subject: [PATCH] warp-server: Add unit tests for websocket routes

---
 warp-server/src/main.rs | 48 +++++++++++++++++++++++++++++++++++++++++
 1 file changed, 48 insertions(+)

diff --git a/warp-server/src/main.rs b/warp-server/src/main.rs
index 7aa6f97..fe36ca7 100644
--- a/warp-server/src/main.rs
+++ b/warp-server/src/main.rs
@@ -98,4 +98,52 @@ mod tests {
         assert_eq!(res.status(), 200);
         assert_eq!(res.body(), "Hello, rust from TestBrowser 0.1!");
     }
+
+    #[tokio::test]
+    async fn test_ws_echo() {
+        let mut client = warp::test::ws()
+            .path("/ws-echo")
+            .handshake(super::api())
+            .await
+            .expect("handshake failed");
+
+        // text
+        client.send_text("Hello").await;
+        let reply = client.recv().await.unwrap();
+        assert_eq!(reply.to_str().unwrap(), "Hello");
+
+        // binary
+        let msg: Vec<u8> = vec![42, 99];
+        client.send(warp::ws::Message::binary(msg.clone())).await;
+        let reply = client.recv().await.unwrap();
+        assert_eq!(reply.as_bytes(), &msg);
+
+        //close
+        client.send(warp::ws::Message::close()).await;
+        client.recv_closed().await.unwrap();
+    }
+
+    #[tokio::test]
+    async fn test_ws_rev() {
+        let mut client = warp::test::ws()
+            .path("/ws-rev")
+            .handshake(super::api())
+            .await
+            .expect("handshake failed");
+
+        // text
+        client.send_text("Hello\n").await;
+        let reply = client.recv().await.unwrap();
+        assert_eq!(reply.to_str().unwrap(), "\nolleH");
+
+        // binary
+        let msg: Vec<u8> = vec![42, 99];
+        client.send(warp::ws::Message::binary(msg)).await;
+        let reply = client.recv().await.unwrap();
+        assert_eq!(reply.as_bytes(), vec![99, 42]);
+
+        //close
+        client.send(warp::ws::Message::close()).await;
+        client.recv_closed().await.unwrap();
+    }
 }
-- 
2.39.5