tests: ws
This commit is contained in:
chesedo
Felipe Noronha
parent
96a398de85
commit
20dbf00931
@ -28,6 +28,7 @@ hyper = { version = "0.14.18", features = ["server"] }
|
|||||||
tokio = { version = "1.17.0", features = ["full"] }
|
tokio = { version = "1.17.0", features = ["full"] }
|
||||||
futures = "0.3.21"
|
futures = "0.3.21"
|
||||||
async-trait = "0.1.53"
|
async-trait = "0.1.53"
|
||||||
|
async-tungstenite = { version = "0.17", features = ["tokio-runtime"] }
|
||||||
tokio-test = "0.4.2"
|
tokio-test = "0.4.2"
|
||||||
test-context = "0.1.3"
|
test-context = "0.1.3"
|
||||||
tokiotest-httpserver = "0.2.1"
|
tokiotest-httpserver = "0.2.1"
|
||||||
@ -38,3 +39,5 @@ hyper-trust-dns = { version = "0.4.2", features = [
|
|||||||
"rustls-webpki"
|
"rustls-webpki"
|
||||||
] }
|
] }
|
||||||
rand = "0.8.5"
|
rand = "0.8.5"
|
||||||
|
tungstenite = "0.17"
|
||||||
|
url = "2.2"
|
||||||
|
|||||||
124
tests/test_websocket.rs
Normal file
124
tests/test_websocket.rs
Normal file
@ -0,0 +1,124 @@
|
|||||||
|
use std::{
|
||||||
|
convert::Infallible,
|
||||||
|
net::{IpAddr, SocketAddr},
|
||||||
|
process::exit,
|
||||||
|
time::Duration,
|
||||||
|
};
|
||||||
|
|
||||||
|
use async_tungstenite::tokio::{accept_async, connect_async};
|
||||||
|
use futures::{SinkExt, StreamExt};
|
||||||
|
use hyper::{
|
||||||
|
client::{connect::dns::GaiResolver, HttpConnector},
|
||||||
|
server::conn::AddrStream,
|
||||||
|
service::{make_service_fn, service_fn},
|
||||||
|
Body, Request, Response, Server,
|
||||||
|
};
|
||||||
|
use hyper_reverse_proxy::ReverseProxy;
|
||||||
|
use test_context::{test_context, AsyncTestContext};
|
||||||
|
use tokio::{net::TcpListener, sync::oneshot::Sender, task::JoinHandle};
|
||||||
|
use tokiotest_httpserver::take_port;
|
||||||
|
use tungstenite::Message;
|
||||||
|
use url::Url;
|
||||||
|
|
||||||
|
lazy_static::lazy_static! {
|
||||||
|
static ref PROXY_CLIENT: ReverseProxy<HttpConnector<GaiResolver>> = {
|
||||||
|
ReverseProxy::new(
|
||||||
|
hyper::Client::new(),
|
||||||
|
)
|
||||||
|
};
|
||||||
|
}
|
||||||
|
|
||||||
|
struct ProxyTestContext {
|
||||||
|
sender: Sender<()>,
|
||||||
|
proxy_handler: JoinHandle<Result<(), hyper::Error>>,
|
||||||
|
ws_handler: JoinHandle<()>,
|
||||||
|
port: u16,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test_context(ProxyTestContext)]
|
||||||
|
#[tokio::test]
|
||||||
|
async fn test_websocket(ctx: &mut ProxyTestContext) {
|
||||||
|
println!("making client connection");
|
||||||
|
let (mut client, _) =
|
||||||
|
connect_async(Url::parse(&format!("ws://127.0.0.1:{}", ctx.port)).unwrap())
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
|
||||||
|
println!("made client connection");
|
||||||
|
client.send(Message::Ping("ping".into())).await.unwrap();
|
||||||
|
let msg = client.next().await.unwrap().unwrap();
|
||||||
|
|
||||||
|
assert!(matches!(msg, Message::Pong(inner) if inner == "pong".as_bytes()));
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn handle(
|
||||||
|
client_ip: IpAddr,
|
||||||
|
req: Request<Body>,
|
||||||
|
backend_port: u16,
|
||||||
|
) -> Result<Response<Body>, Infallible> {
|
||||||
|
match PROXY_CLIENT
|
||||||
|
.call(
|
||||||
|
client_ip,
|
||||||
|
format!("http://127.0.0.1:{}", backend_port).as_str(),
|
||||||
|
req,
|
||||||
|
)
|
||||||
|
.await
|
||||||
|
{
|
||||||
|
Ok(response) => Ok(response),
|
||||||
|
Err(err) => panic!("did not expect error: {:?}", err),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[async_trait::async_trait]
|
||||||
|
impl<'a> AsyncTestContext for ProxyTestContext {
|
||||||
|
async fn setup() -> ProxyTestContext {
|
||||||
|
tokio::spawn(async {
|
||||||
|
tokio::time::sleep(Duration::from_secs(5)).await;
|
||||||
|
println!("Unit test executed too long, perhaps its stuck...");
|
||||||
|
exit(1);
|
||||||
|
});
|
||||||
|
|
||||||
|
let (sender, receiver) = tokio::sync::oneshot::channel::<()>();
|
||||||
|
let ws_port = take_port();
|
||||||
|
|
||||||
|
let ws_handler = tokio::spawn(async move {
|
||||||
|
let ws_server = TcpListener::bind(("127.0.0.1", ws_port)).await.unwrap();
|
||||||
|
|
||||||
|
while let Ok((stream, addr)) = ws_server.accept().await {
|
||||||
|
println!("incoming connection: {addr}");
|
||||||
|
let mut websocket = accept_async(stream).await.unwrap();
|
||||||
|
|
||||||
|
let msg = websocket.next().await.unwrap().unwrap();
|
||||||
|
assert!(matches!(msg, Message::Ping(inner) if inner == "ping".as_bytes()));
|
||||||
|
println!("past ping");
|
||||||
|
|
||||||
|
websocket.send(Message::Pong("pong".into())).await.unwrap();
|
||||||
|
println!("past pong");
|
||||||
|
}
|
||||||
|
});
|
||||||
|
|
||||||
|
let make_svc = make_service_fn(move |conn: &AddrStream| {
|
||||||
|
let remote_addr = conn.remote_addr().ip();
|
||||||
|
async move { Ok::<_, Infallible>(service_fn(move |req| handle(remote_addr, req, ws_port))) }
|
||||||
|
});
|
||||||
|
|
||||||
|
let port = take_port();
|
||||||
|
let addr = SocketAddr::new("127.0.0.1".parse().unwrap(), port);
|
||||||
|
let server = Server::bind(&addr)
|
||||||
|
.serve(make_svc)
|
||||||
|
.with_graceful_shutdown(async {
|
||||||
|
receiver.await.ok();
|
||||||
|
});
|
||||||
|
let proxy_handler = tokio::spawn(server);
|
||||||
|
ProxyTestContext {
|
||||||
|
sender,
|
||||||
|
proxy_handler,
|
||||||
|
ws_handler,
|
||||||
|
port,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
async fn teardown(self) {
|
||||||
|
let _ = self.sender.send(()).unwrap();
|
||||||
|
let _ = tokio::join!(self.proxy_handler, self.ws_handler);
|
||||||
|
}
|
||||||
|
}
|
||||||
Loading…
Reference in New Issue
Block a user