Add websocket integration test (#2477)

axum/Cargo.toml

@@ -124,6 +124,7 @@ serde_json = "1.0"
 time = { version = "0.3", features = ["serde-human-readable"] }
 tokio = { package = "tokio", version = "1.25.0", features = ["macros", "rt", "rt-multi-thread", "net", "test-util"] }
 tokio-stream = "0.1"
+tokio-tungstenite = "0.21"
 tracing = "0.1"
 tracing-subscriber = { version = "0.3", features = ["json"] }
 uuid = { version = "1.0", features = ["serde", "v4"] }

axum/src/extract/ws.rs

@@ -830,9 +830,12 @@ pub mod close_code {
 
 #[cfg(test)]
 mod tests {
+    use std::future::ready;
+
     use super::*;
-    use crate::{body::Body, routing::get, Router};
+    use crate::{body::Body, routing::get, test_helpers::spawn_service, Router};
     use http::{Request, Version};
+    use tokio_tungstenite::tungstenite;
     use tower::ServiceExt;
 
     #[crate::test]
@@ -877,4 +880,47 @@ mod tests {
         }
         let _: Router = Router::new().route("/", get(handler));
     }
+
+    #[crate::test]
+    async fn integration_test() {
+        let app = Router::new().route(
+            "/echo",
+            get(|ws: WebSocketUpgrade| ready(ws.on_upgrade(handle_socket))),
+        );
+
+        async fn handle_socket(mut socket: WebSocket) {
+            while let Some(Ok(msg)) = socket.recv().await {
+                match msg {
+                    Message::Text(_) | Message::Binary(_) | Message::Close(_) => {
+                        if socket.send(msg).await.is_err() {
+                            break;
+                        }
+                    }
+                    Message::Ping(_) | Message::Pong(_) => {
+                        // tungstenite will respond to pings automatically
+                    }
+                }
+            }
+        }
+
+        let addr = spawn_service(app);
+        let (mut socket, _response) = tokio_tungstenite::connect_async(format!("ws://{addr}/echo"))
+            .await
+            .unwrap();
+
+        let input = tungstenite::Message::Text("foobar".to_owned());
+        socket.send(input.clone()).await.unwrap();
+        let output = socket.next().await.unwrap().unwrap();
+        assert_eq!(input, output);
+
+        socket
+            .send(tungstenite::Message::Ping("ping".to_owned().into_bytes()))
+            .await
+            .unwrap();
+        let output = socket.next().await.unwrap().unwrap();
+        assert_eq!(
+            output,
+            tungstenite::Message::Pong("ping".to_owned().into_bytes())
+        );
+    }
 }

axum/src/test_helpers/test_client.rs

@@ -10,13 +10,7 @@ use tokio::net::TcpListener;
 use tower::make::Shared;
 use tower_service::Service;
 
-pub(crate) struct TestClient {
-    client: reqwest::Client,
-    addr: SocketAddr,
-}
-
-impl TestClient {
-    pub(crate) fn new<S>(svc: S) -> Self
+pub(crate) fn spawn_service<S>(svc: S) -> SocketAddr
 where
     S: Service<Request, Response = Response, Error = Infallible> + Clone + Send + 'static,
     S::Future: Send,
@@ -34,6 +28,22 @@ impl TestClient {
             .expect("server error")
     });
 
+    addr
+}
+
+pub(crate) struct TestClient {
+    client: reqwest::Client,
+    addr: SocketAddr,
+}
+
+impl TestClient {
+    pub(crate) fn new<S>(svc: S) -> Self
+    where
+        S: Service<Request, Response = Response, Error = Infallible> + Clone + Send + 'static,
+        S::Future: Send,
+    {
+        let addr = spawn_service(svc);
+
         let client = reqwest::Client::builder()
             .redirect(reqwest::redirect::Policy::none())
             .build()