ws_mock/
utils.rs

1//! Utilities used throughout testing and implementations that may be of use.
2use crate::ws_mock_server::WsMockServer;
3use futures_util::stream::SplitStream;
4use futures_util::{SinkExt, StreamExt};
5use std::time::Duration;
6use tokio::net::TcpStream;
7use tokio_tungstenite::tungstenite::Message;
8use tokio_tungstenite::{MaybeTlsStream, WebSocketStream, connect_async};
9
10/// Send a single message to the server, and return the `SplitStream` receiver created by
11/// connecting to the server.
12///
13/// Useful for one-off messages and optionally reading responses afterwards.
14pub async fn send_to_server(
15    server: &WsMockServer,
16    message: String,
17) -> SplitStream<WebSocketStream<MaybeTlsStream<TcpStream>>> {
18    let (stream, _resp) = connect_async(server.uri().await)
19        .await
20        .expect("Connecting failed");
21
22    let (mut send, recv) = stream.split();
23
24    send.send(Message::from(message)).await.unwrap();
25
26    recv
27}
28
29/// Given a `SplitStream` receiver, receive messages until timing out and return all messages in a `Vec<String>`.
30///
31/// Useful for receiving a batch of messages and later making assertions about them.
32pub async fn collect_all_messages(
33    mut ws_recv: SplitStream<WebSocketStream<MaybeTlsStream<TcpStream>>>,
34    timeout: Duration,
35) -> Vec<Message> {
36    let mut received = Vec::new();
37
38    while let Ok(Some(Ok(message))) = tokio::time::timeout(timeout, ws_recv.next()).await {
39        received.push(message);
40    }
41    received
42}