234 lines
14 KiB
Rust
234 lines
14 KiB
Rust
// file: crates/ksp-onchain-transport-lib/unit_tests/ws_cluster.rs
|
|
// version: 2
|
|
|
|
use futures_util::SinkExt; // rust-rules: trait-import
|
|
use futures_util::StreamExt; // rust-rules: trait-import
|
|
|
|
fn local_endpoint(url: &str) -> crate::WsEndpointSettings {
|
|
return crate::WsEndpointSettings::new(
|
|
"local_ws_cluster",
|
|
true,
|
|
crate::WsProviderName::new("local-fixture"),
|
|
crate::WsClusterName::new("local"),
|
|
crate::WsProtocolKind::SolanaStandard,
|
|
crate::WsEndpointUrl::parse(url).expect("local test WebSocket URL must parse"),
|
|
crate::WsSessionSettings::default(),
|
|
);
|
|
}
|
|
|
|
async fn bind_local_listener() -> (tokio::net::TcpListener, std::string::String) {
|
|
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.expect("local listener must bind");
|
|
let address = listener.local_addr().expect("local listener must expose address");
|
|
return (listener, format!("ws://{address}"));
|
|
}
|
|
|
|
async fn read_request(websocket: &mut tokio_tungstenite::WebSocketStream<tokio::net::TcpStream>) -> serde_json::Value {
|
|
let message = websocket.next().await.expect("request message must exist").expect("request message must decode");
|
|
let text = message.to_text().expect("request must be text");
|
|
return serde_json::from_str(text).expect("request must contain JSON");
|
|
}
|
|
|
|
async fn send_result(websocket: &mut tokio_tungstenite::WebSocketStream<tokio::net::TcpStream>, request: &serde_json::Value, result: serde_json::Value) {
|
|
let id = request.get("id").and_then(serde_json::Value::as_u64).expect("request id must be numeric");
|
|
let response = serde_json::json!({"jsonrpc":"2.0","id":id,"result":result});
|
|
websocket.send(tokio_tungstenite::tungstenite::Message::Text(response.to_string().into())).await.expect("local response must send");
|
|
}
|
|
|
|
async fn send_notification(websocket: &mut tokio_tungstenite::WebSocketStream<tokio::net::TcpStream>, method: &str, remote_id: u64, result: serde_json::Value) {
|
|
let notification = serde_json::json!({"jsonrpc":"2.0","method":method,"params":{"result":result,"subscription":remote_id}});
|
|
websocket.send(tokio_tungstenite::tungstenite::Message::Text(notification.to_string().into())).await.expect("local notification must send");
|
|
}
|
|
|
|
async fn wait_for_close_frame(websocket: &mut tokio_tungstenite::WebSocketStream<tokio::net::TcpStream>) {
|
|
loop {
|
|
let message = websocket.next().await;
|
|
match message {
|
|
std::option::Option::Some(std::result::Result::Ok(tokio_tungstenite::tungstenite::Message::Close(_))) => return,
|
|
std::option::Option::Some(std::result::Result::Ok(_)) => {},
|
|
std::option::Option::Some(std::result::Result::Err(_)) | std::option::Option::None => return,
|
|
}
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn slot_notification_decoder_preserves_slot_parent_and_root() {
|
|
let notification =
|
|
super::decode_slot_notification("slotSubscribe", serde_json::json!({"slot":76,"parent":75,"root":44})).expect("slot notification must decode");
|
|
assert_eq!(notification.slot(), 76);
|
|
assert_eq!(notification.parent(), 75);
|
|
assert_eq!(notification.root(), 44);
|
|
assert!(super::decode_slot_notification("slotSubscribe", serde_json::json!({"slot":76,"parent":75})).is_err());
|
|
}
|
|
|
|
#[tokio::test(flavor = "current_thread")]
|
|
async fn stable_slot_and_root_wrappers_use_no_params_decode_exact_notifications_and_unsubscribe_by_handle() {
|
|
let (listener, url) = bind_local_listener().await;
|
|
let server = tokio::spawn(async move {
|
|
let (stream, _) = listener.accept().await.expect("local server must accept client");
|
|
let mut websocket = tokio_tungstenite::accept_async(stream).await.expect("local WebSocket handshake must succeed");
|
|
let slot_subscribe = read_request(&mut websocket).await;
|
|
assert_eq!(slot_subscribe["method"], serde_json::json!("slotSubscribe"));
|
|
assert_eq!(slot_subscribe["params"], serde_json::json!([]));
|
|
send_result(&mut websocket, &slot_subscribe, serde_json::json!(201)).await;
|
|
send_notification(&mut websocket, "slotNotification", 201, serde_json::json!({"slot":76,"parent":75,"root":44})).await;
|
|
let root_subscribe = read_request(&mut websocket).await;
|
|
assert_eq!(root_subscribe["method"], serde_json::json!("rootSubscribe"));
|
|
assert_eq!(root_subscribe["params"], serde_json::json!([]));
|
|
send_result(&mut websocket, &root_subscribe, serde_json::json!(202)).await;
|
|
send_notification(&mut websocket, "rootNotification", 202, serde_json::json!(42)).await;
|
|
let slot_unsubscribe = read_request(&mut websocket).await;
|
|
assert_eq!(slot_unsubscribe["method"], serde_json::json!("slotUnsubscribe"));
|
|
assert_eq!(slot_unsubscribe["params"], serde_json::json!([201]));
|
|
send_result(&mut websocket, &slot_unsubscribe, serde_json::json!(true)).await;
|
|
let root_unsubscribe = read_request(&mut websocket).await;
|
|
assert_eq!(root_unsubscribe["method"], serde_json::json!("rootUnsubscribe"));
|
|
assert_eq!(root_unsubscribe["params"], serde_json::json!([202]));
|
|
send_result(&mut websocket, &root_unsubscribe, serde_json::json!(true)).await;
|
|
wait_for_close_frame(&mut websocket).await;
|
|
});
|
|
let session = crate::WsSession::connect(local_endpoint(url.as_str())).await.expect("client handshake must succeed");
|
|
let mut slot_subscription = session.slot_subscribe().await.expect("slotSubscribe must register");
|
|
let slot = slot_subscription.recv().await.expect("slot notification must arrive").expect("slot notification must decode");
|
|
assert_eq!((slot.slot(), slot.parent(), slot.root()), (76, 75, 44));
|
|
let mut root_subscription = session.root_subscribe().await.expect("rootSubscribe must register");
|
|
let root = root_subscription.recv().await.expect("root notification must arrive").expect("root notification must decode");
|
|
assert_eq!(root, 42);
|
|
assert!(slot_subscription.unsubscribe().await.expect("slot unsubscribe must complete"));
|
|
assert!(root_subscription.unsubscribe().await.expect("root unsubscribe must complete"));
|
|
session.close().await.expect("session close must complete");
|
|
server.await.expect("local server task must complete");
|
|
}
|
|
|
|
#[test]
|
|
fn slots_update_decoder_preserves_all_known_variants_and_unknown_bounded_fallback() {
|
|
let cases = [
|
|
(serde_json::json!({"slot":1,"timestamp":10,"type":"firstShredReceived"}), "firstShredReceived"),
|
|
(serde_json::json!({"slot":2,"timestamp":20,"type":"completed"}), "completed"),
|
|
(serde_json::json!({"slot":3,"timestamp":30,"type":"createdBank","parent":2}), "createdBank"),
|
|
(
|
|
serde_json::json!({
|
|
"slot":4,
|
|
"timestamp":40,
|
|
"type":"frozen",
|
|
"stats":{"maxTransactionsPerEntry":64,"numFailedTransactions":1,"numSuccessfulTransactions":9,"numTransactionEntries":3}
|
|
}),
|
|
"frozen",
|
|
),
|
|
(serde_json::json!({"slot":5,"timestamp":50,"type":"dead","err":"fixture dead"}), "dead"),
|
|
(serde_json::json!({"slot":6,"timestamp":60,"type":"optimisticConfirmation"}), "optimisticConfirmation"),
|
|
(serde_json::json!({"slot":7,"timestamp":70,"type":"root"}), "root"),
|
|
];
|
|
for (wire, expected_type) in cases {
|
|
let update = super::decode_slots_update_notification("slotsUpdatesSubscribe", wire).expect("known slot update must decode");
|
|
assert_eq!(update.update_type(), expected_type);
|
|
assert!(update.slot().is_some());
|
|
assert!(update.timestamp().is_some());
|
|
assert!(update.unknown_raw().is_none());
|
|
}
|
|
let frozen = super::decode_slots_update_notification(
|
|
"slotsUpdatesSubscribe",
|
|
serde_json::json!({
|
|
"slot":4,
|
|
"timestamp":40,
|
|
"type":"frozen",
|
|
"stats":{"maxTransactionsPerEntry":64,"numFailedTransactions":1,"numSuccessfulTransactions":9,"numTransactionEntries":3}
|
|
}),
|
|
)
|
|
.expect("frozen update must decode");
|
|
match frozen {
|
|
crate::SolanaSlotUpdate::Frozen { stats, .. } => {
|
|
assert_eq!(stats.max_transactions_per_entry(), 64);
|
|
assert_eq!(stats.num_failed_transactions(), 1);
|
|
assert_eq!(stats.num_successful_transactions(), 9);
|
|
assert_eq!(stats.num_transaction_entries(), 3);
|
|
},
|
|
_ => panic!("fixture must decode as frozen"),
|
|
}
|
|
let unknown_wire = serde_json::json!({"slot":8,"timestamp":80,"type":"futureBankState","futureField":{"x":1}});
|
|
let unknown = super::decode_slots_update_notification("slotsUpdatesSubscribe", unknown_wire.clone()).expect("unknown update must remain consumable");
|
|
assert_eq!(unknown.update_type(), "futureBankState");
|
|
assert_eq!(unknown.slot(), std::option::Option::Some(8));
|
|
assert_eq!(unknown.timestamp(), std::option::Option::Some(80));
|
|
assert_eq!(unknown.unknown_raw(), std::option::Option::Some(&unknown_wire));
|
|
}
|
|
|
|
#[test]
|
|
fn slots_update_known_variants_require_their_variant_specific_fields() {
|
|
assert!(super::decode_slots_update_notification("slotsUpdatesSubscribe", serde_json::json!({"slot":3,"timestamp":30,"type":"createdBank"})).is_err());
|
|
assert!(super::decode_slots_update_notification("slotsUpdatesSubscribe", serde_json::json!({"slot":4,"timestamp":40,"type":"frozen"})).is_err());
|
|
assert!(super::decode_slots_update_notification("slotsUpdatesSubscribe", serde_json::json!({"slot":5,"timestamp":50,"type":"dead"})).is_err());
|
|
}
|
|
|
|
#[test]
|
|
fn vote_notification_decoder_preserves_timestamp_omitted_null_and_value() {
|
|
let pubkey = "11111111111111111111111111111111";
|
|
for (wire, expected_timestamp) in [
|
|
(serde_json::json!({"votePubkey":pubkey,"slots":[1,2],"hash":"hash-a","signature":"sig-a"}), std::option::Option::None),
|
|
(serde_json::json!({"votePubkey":pubkey,"slots":[1,2],"hash":"hash-b","timestamp":null,"signature":"sig-b"}), std::option::Option::None),
|
|
(serde_json::json!({"votePubkey":pubkey,"slots":[1,2],"hash":"hash-c","timestamp":123,"signature":"sig-c"}), std::option::Option::Some(123)),
|
|
] {
|
|
let vote = super::decode_vote_notification("voteSubscribe", wire).expect("vote notification must decode");
|
|
assert_eq!(vote.vote_pubkey().to_string(), pubkey);
|
|
assert_eq!(vote.slots(), &[1, 2]);
|
|
assert_eq!(vote.timestamp(), expected_timestamp);
|
|
}
|
|
}
|
|
|
|
#[tokio::test(flavor = "current_thread")]
|
|
async fn unstable_slots_updates_and_vote_wrappers_use_no_params_and_handle_unsubscribe() {
|
|
let (listener, url) = bind_local_listener().await;
|
|
let server = tokio::spawn(async move {
|
|
let (stream, _) = listener.accept().await.expect("local server must accept client");
|
|
let mut websocket = tokio_tungstenite::accept_async(stream).await.expect("local WebSocket handshake must succeed");
|
|
let slots_subscribe = read_request(&mut websocket).await;
|
|
assert_eq!(slots_subscribe["method"], serde_json::json!("slotsUpdatesSubscribe"));
|
|
assert_eq!(slots_subscribe["params"], serde_json::json!([]));
|
|
send_result(&mut websocket, &slots_subscribe, serde_json::json!(401)).await;
|
|
send_notification(
|
|
&mut websocket,
|
|
"slotsUpdatesNotification",
|
|
401,
|
|
serde_json::json!({"slot":76,"timestamp":1625081266243_i64,"type":"optimisticConfirmation"}),
|
|
)
|
|
.await;
|
|
let vote_subscribe = read_request(&mut websocket).await;
|
|
assert_eq!(vote_subscribe["method"], serde_json::json!("voteSubscribe"));
|
|
assert_eq!(vote_subscribe["params"], serde_json::json!([]));
|
|
send_result(&mut websocket, &vote_subscribe, serde_json::json!(402)).await;
|
|
send_notification(
|
|
&mut websocket,
|
|
"voteNotification",
|
|
402,
|
|
serde_json::json!({
|
|
"votePubkey":"11111111111111111111111111111111",
|
|
"slots":[75,76],
|
|
"hash":"fixture-hash",
|
|
"timestamp":null,
|
|
"signature":"fixture-signature"
|
|
}),
|
|
)
|
|
.await;
|
|
let slots_unsubscribe = read_request(&mut websocket).await;
|
|
assert_eq!(slots_unsubscribe["method"], serde_json::json!("slotsUpdatesUnsubscribe"));
|
|
assert_eq!(slots_unsubscribe["params"], serde_json::json!([401]));
|
|
send_result(&mut websocket, &slots_unsubscribe, serde_json::json!(true)).await;
|
|
let vote_unsubscribe = read_request(&mut websocket).await;
|
|
assert_eq!(vote_unsubscribe["method"], serde_json::json!("voteUnsubscribe"));
|
|
assert_eq!(vote_unsubscribe["params"], serde_json::json!([402]));
|
|
send_result(&mut websocket, &vote_unsubscribe, serde_json::json!(true)).await;
|
|
wait_for_close_frame(&mut websocket).await;
|
|
});
|
|
let session = crate::WsSession::connect(local_endpoint(url.as_str())).await.expect("client handshake must succeed");
|
|
let mut slots_subscription = session.slots_updates_subscribe().await.expect("slotsUpdatesSubscribe must register");
|
|
let update = slots_subscription.recv().await.expect("slots update must arrive").expect("slots update must decode");
|
|
assert_eq!(update.update_type(), "optimisticConfirmation");
|
|
let mut vote_subscription = session.vote_subscribe().await.expect("voteSubscribe must register");
|
|
let vote = vote_subscription.recv().await.expect("vote notification must arrive").expect("vote notification must decode");
|
|
assert_eq!(vote.slots(), &[75, 76]);
|
|
assert!(vote.timestamp().is_none());
|
|
assert!(slots_subscription.unsubscribe().await.expect("slots update unsubscribe must complete"));
|
|
assert!(vote_subscription.unsubscribe().await.expect("vote unsubscribe must complete"));
|
|
session.close().await.expect("session close must complete");
|
|
server.await.expect("local server task must complete");
|
|
}
|