warp-server: Add unit tests for websocket routes

[learn-rust.git] / warp-server / src / main.rs

diff --git a/warp-server/src/main.rs b/warp-server/src/main.rs
index 06977e62cf4e50911b5e00ba6f866c2daa49dc52..fe36ca7faf989572f7eb765e64a4e8497fde8f97 100644 (file)
--- a/warp-server/src/main.rs
+++ b/warp-server/src/main.rs
@@ -1,4 +1,4 @@
-use futures_util::{FutureExt, StreamExt};
+use futures_util::{FutureExt, StreamExt, SinkExt};
 use warp::Filter;
 
 // GET /hello/warp => 200 OK with body "Hello, warp!"
@@ -17,25 +17,133 @@ async fn ws_echo_connected(websocket: warp::ws::WebSocket) {
     }).await;
 }
 
-#[tokio::main]
-async fn main() {
-    env_logger::Builder::from_env(env_logger::Env::default().default_filter_or("info")).init();
+// websocat ws://127.0.0.1:3030/ws-rev
+async fn ws_rev_connected(websocket: warp::ws::WebSocket) {
+    // echo all messages back
+    tokio::task::spawn(async {
+        let (mut tx, mut rx) = websocket.split();
+        while let Some(message) = rx.next().await {
+            let msg = match message {
+                Ok(msg) =>  msg,
+                Err(e) => {
+                    log::error!("websocket error: {}", e);
+                    break;
+                }
+            };
+            log::info!("ws_rev_connected: got message: {:?}", msg);
+
+            if msg.is_close() {
+                break;
+            }
+            if msg.is_text() {
+                let text = msg.to_str().unwrap();
+                let rev = text.chars().rev().collect::<String>();
+                if let Err(e) = tx.send(warp::ws::Message::text(rev)).await {
+                    // disconnected
+                    log::info!("peer disconnected: {}", e);
+                    break;
+                }
+            }
+            if msg.is_binary() {
+                let mut rev = msg.into_bytes();
+                rev.reverse();
+                if let Err(e) = tx.send(warp::ws::Message::binary(rev)).await {
+                    // disconnected
+                    log::info!("peer disconnected: {}", e);
+                    break;
+                }
+            }
+        }
+        log::info!("ws_rev_connected ended");
+    });
+}
 
+pub fn api() -> impl Filter<Extract = impl warp::Reply, Error = warp::Rejection> + Clone {
     let hello = warp::path!("hello" / String)
         .and(warp::header::<String>("user-agent"))
         .and_then(hello);
 
     let ws_echo = warp::path("ws-echo")
         .and(warp::ws())
-        .map(|ws: warp::ws::Ws| {
-            ws.on_upgrade(|websocket| { ws_echo_connected(websocket) })
-        });
+        .map(|ws: warp::ws::Ws| { ws.on_upgrade(ws_echo_connected) });
+
+    let ws_rev = warp::path("ws-rev")
+        .and(warp::ws())
+        .map(|ws: warp::ws::Ws| { ws.on_upgrade(ws_rev_connected) });
 
-    let api = hello
+    hello
         .or(ws_echo)
-        .with(warp::log("warp-server"));
+        .or(ws_rev)
+        .with(warp::log("warp-server"))
+}
+
+#[tokio::main]
+async fn main() {
+    env_logger::Builder::from_env(env_logger::Env::default().default_filter_or("info")).init();
 
-    warp::serve(api)
+    warp::serve(api())
         .run(([127, 0, 0, 1], 3030))
         .await;
 }
+
+#[cfg(test)]
+mod tests {
+    #[tokio::test]
+    async fn test_hello() {
+        let res = warp::test::request()
+            .path("/hello/rust")
+            .header("user-agent", "TestBrowser 0.1")
+            .reply(&super::api())
+            .await;
+        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();
+    }
+}