#[cfg(feature = "websocket")]
use pmcp::server::transport::websocket::WebSocketServerBuilder;
#[cfg(feature = "websocket")]
use pmcp::shared::{Transport, TransportMessage};
#[cfg(feature = "websocket")]
use pmcp::types::{ClientNotification, Notification, ProgressNotification};
#[cfg(feature = "websocket")]
use std::time::Duration;
#[cfg(feature = "websocket")]
use tokio::time::timeout;
#[cfg(feature = "websocket")]
#[tokio::test]
async fn test_websocket_server_bind() {
let mut transport = WebSocketServerBuilder::new()
.bind_addr("127.0.0.1:0".parse().unwrap())
.build();
transport.bind().await.expect("Failed to bind");
assert!(!transport.is_connected());
}
#[cfg(feature = "websocket")]
#[tokio::test]
async fn test_websocket_server_accept_timeout() {
let mut transport = WebSocketServerBuilder::new()
.bind_addr("127.0.0.1:0".parse().unwrap())
.build();
transport.bind().await.expect("Failed to bind");
let result = timeout(Duration::from_millis(100), transport.accept()).await;
assert!(result.is_err(), "Accept should timeout");
}
#[cfg(feature = "websocket")]
#[tokio::test]
async fn test_websocket_server_client_connection() {
use tokio_tungstenite::connect_async;
let mut server_transport = WebSocketServerBuilder::new()
.bind_addr("127.0.0.1:9003".parse().unwrap())
.build();
server_transport.bind().await.expect("Failed to bind");
let server_handle = tokio::spawn(async move {
server_transport.accept().await.expect("Failed to accept");
assert!(server_transport.is_connected());
server_transport
});
tokio::time::sleep(Duration::from_millis(50)).await;
let (ws_stream, _) = connect_async("ws://127.0.0.1:9003")
.await
.expect("Failed to connect");
let server_transport = server_handle.await.unwrap();
assert!(server_transport.is_connected());
drop(ws_stream);
}
#[cfg(feature = "websocket")]
#[tokio::test]
async fn test_websocket_server_send_receive() {
use futures::{SinkExt, StreamExt};
use tokio_tungstenite::{connect_async, tungstenite::Message};
let mut server_transport = WebSocketServerBuilder::new()
.bind_addr("127.0.0.1:9004".parse().unwrap())
.build();
server_transport.bind().await.expect("Failed to bind");
let server_handle = tokio::spawn(async move {
server_transport.accept().await.expect("Failed to accept");
let notification = TransportMessage::Notification(Notification::Client(
ClientNotification::Progress(ProgressNotification::new(
pmcp::types::ProgressToken::String("test-token".to_string()),
50.0,
Some("Testing".to_string()),
)),
));
server_transport
.send(notification)
.await
.expect("Failed to send");
let received = server_transport.receive().await.expect("Failed to receive");
(server_transport, received)
});
tokio::time::sleep(Duration::from_millis(50)).await;
let (ws_stream, _) = connect_async("ws://127.0.0.1:9004")
.await
.expect("Failed to connect");
let (mut ws_sender, mut ws_receiver) = ws_stream.split();
let msg = ws_receiver.next().await.unwrap().unwrap();
if let Message::Text(text) = msg {
assert!(text.contains("progress"));
assert!(text.contains("Testing"));
} else {
panic!("Expected text message");
}
let request_msg = serde_json::json!({
"jsonrpc": "2.0",
"method": "ping",
"id": "test-ping"
});
ws_sender
.send(Message::Text(request_msg.to_string().into()))
.await
.unwrap();
let (server_transport, received) = server_handle.await.unwrap();
match received {
TransportMessage::Request { id, request: _ } => {
assert_eq!(id, pmcp::types::RequestId::from("test-ping"));
},
_ => panic!("Expected request message, got: {:?}", received),
}
drop(server_transport);
}
#[cfg(feature = "websocket")]
#[tokio::test]
async fn test_websocket_server_close() {
let mut transport = WebSocketServerBuilder::new()
.bind_addr("127.0.0.1:9005".parse().unwrap())
.build();
transport.bind().await.expect("Failed to bind");
transport.close().await.expect("Failed to close");
assert!(!transport.is_connected());
}
#[cfg(feature = "websocket")]
#[tokio::test]
async fn test_websocket_server_builder_options() {
let transport = WebSocketServerBuilder::new()
.bind_addr("127.0.0.1:9006".parse().unwrap())
.max_frame_size(1024 * 1024)
.max_message_size(2 * 1024 * 1024)
.accept_unmasked_frames(true)
.build();
assert_eq!(transport.transport_type(), "websocket-server");
}