use compio_net::{TcpListener, TcpStream};
use compio_ws::{accept_async, client_async};
use futures_channel::oneshot;
use tungstenite::Message;
#[compio_macros::test]
async fn test_handshake() {
let (tx, rx) = oneshot::channel();
compio_runtime::spawn(async move {
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = listener.local_addr().unwrap();
tx.send(addr).unwrap();
let (stream, _) = listener.accept().await.unwrap();
let _ws = accept_async(stream).await.expect("Server handshake failed");
})
.detach();
let addr = rx.await.expect("Failed to wait for server");
let tcp = TcpStream::connect(&addr).await.expect("Failed to connect");
let (_ws, _response) = client_async(&format!("ws://{}", addr), tcp)
.await
.expect("Client handshake failed");
}
#[compio_macros::test]
async fn test_echo_message() {
let (tx, rx) = oneshot::channel();
compio_runtime::spawn(async move {
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = listener.local_addr().unwrap();
tx.send(addr).unwrap();
let (stream, _) = listener.accept().await.unwrap();
let mut ws = accept_async(stream).await.unwrap();
let msg = ws.read().await.unwrap();
ws.send(msg).await.unwrap();
})
.detach();
let addr = rx.await.unwrap();
let tcp = TcpStream::connect(&addr).await.unwrap();
let (mut ws, _) = client_async(&format!("ws://{}", addr), tcp).await.unwrap();
let test_msg = "Hello, WebSocket!";
ws.send(Message::Text(test_msg.into())).await.unwrap();
let response = ws.read().await.unwrap();
assert_eq!(response, Message::Text(test_msg.into()));
}
#[compio_macros::test]
async fn test_binary_message() {
let (tx, rx) = oneshot::channel();
compio_runtime::spawn(async move {
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = listener.local_addr().unwrap();
tx.send(addr).unwrap();
let (stream, _) = listener.accept().await.unwrap();
let mut ws = accept_async(stream).await.unwrap();
let msg = ws.read().await.unwrap();
ws.send(msg).await.unwrap();
})
.detach();
let addr = rx.await.unwrap();
let tcp = TcpStream::connect(&addr).await.unwrap();
let (mut ws, _) = client_async(&format!("ws://{}", addr), tcp).await.unwrap();
let test_data = vec![1, 2, 3, 4, 5];
ws.send(Message::Binary(test_data.clone().into()))
.await
.unwrap();
let response = ws.read().await.unwrap();
assert_eq!(response, Message::Binary(test_data.into()));
}
#[compio_macros::test]
async fn test_ping_pong() {
let (tx, rx) = oneshot::channel();
compio_runtime::spawn(async move {
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = listener.local_addr().unwrap();
tx.send(addr).unwrap();
let (stream, _) = listener.accept().await.unwrap();
let mut ws = accept_async(stream).await.unwrap();
let msg = ws.read().await.unwrap();
if let Message::Ping(data) = msg {
ws.send(Message::Pong(data)).await.unwrap();
}
})
.detach();
let addr = rx.await.unwrap();
let tcp = TcpStream::connect(&addr).await.unwrap();
let (mut ws, _) = client_async(&format!("ws://{}", addr), tcp).await.unwrap();
let ping_data = vec![42];
ws.send(Message::Ping(ping_data.clone().into()))
.await
.unwrap();
let response = ws.read().await.unwrap();
assert_eq!(response, Message::Pong(ping_data.into()));
}
#[compio_macros::test]
async fn test_close_handshake() {
let (tx, rx) = oneshot::channel();
compio_runtime::spawn(async move {
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = listener.local_addr().unwrap();
tx.send(addr).unwrap();
let (stream, _) = listener.accept().await.unwrap();
let mut ws = accept_async(stream).await.unwrap();
let msg = ws.read().await.unwrap();
assert!(msg.is_close());
})
.detach();
let addr = rx.await.unwrap();
let tcp = TcpStream::connect(&addr).await.unwrap();
let (mut ws, _) = client_async(&format!("ws://{}", addr), tcp).await.unwrap();
ws.close(None).await.unwrap();
}
#[compio_macros::test]
async fn test_multiple_messages() {
let (tx, rx) = oneshot::channel();
compio_runtime::spawn(async move {
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = listener.local_addr().unwrap();
tx.send(addr).unwrap();
let (stream, _) = listener.accept().await.unwrap();
let mut ws = accept_async(stream).await.unwrap();
for _ in 0..3 {
let msg = ws.read().await.unwrap();
ws.send(msg).await.unwrap();
}
})
.detach();
let addr = rx.await.unwrap();
let tcp = TcpStream::connect(&addr).await.unwrap();
let (mut ws, _) = client_async(&format!("ws://{}", addr), tcp).await.unwrap();
for i in 0..3 {
let msg = format!("Message {}", i);
ws.send(Message::Text(msg.clone().into())).await.unwrap();
let response = ws.read().await.unwrap();
assert_eq!(response, Message::Text(msg.into()));
}
}