diff --git a/axum/Cargo.toml b/axum/Cargo.toml index da8f5ef2..18893f4f 100644 --- a/axum/Cargo.toml +++ b/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"] } diff --git a/axum/src/extract/ws.rs b/axum/src/extract/ws.rs index 7781d6c2..ecf320ab 100644 --- a/axum/src/extract/ws.rs +++ b/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()) + ); + } } diff --git a/axum/src/test_helpers/test_client.rs b/axum/src/test_helpers/test_client.rs index 89d74fc9..f6751fd2 100644 --- a/axum/src/test_helpers/test_client.rs +++ b/axum/src/test_helpers/test_client.rs @@ -10,6 +10,27 @@ use tokio::net::TcpListener; use tower::make::Shared; use tower_service::Service; +pub(crate) fn spawn_service(svc: S) -> SocketAddr +where + S: Service + Clone + Send + 'static, + S::Future: Send, +{ + let std_listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap(); + std_listener.set_nonblocking(true).unwrap(); + let listener = TcpListener::from_std(std_listener).unwrap(); + + let addr = listener.local_addr().unwrap(); + println!("Listening on {addr}"); + + tokio::spawn(async move { + serve(listener, Shared::new(svc)) + .await + .expect("server error") + }); + + addr +} + pub(crate) struct TestClient { client: reqwest::Client, addr: SocketAddr, @@ -21,18 +42,7 @@ impl TestClient { S: Service + Clone + Send + 'static, S::Future: Send, { - let std_listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap(); - std_listener.set_nonblocking(true).unwrap(); - let listener = TcpListener::from_std(std_listener).unwrap(); - - let addr = listener.local_addr().unwrap(); - println!("Listening on {addr}"); - - tokio::spawn(async move { - serve(listener, Shared::new(svc)) - .await - .expect("server error") - }); + let addr = spawn_service(svc); let client = reqwest::Client::builder() .redirect(reqwest::redirect::Policy::none())