use crate::commands;
use crate::core::{Cmd, Message};
use std::sync::Arc;
use tokio::sync::Mutex;
#[derive(Debug, Clone)]
pub enum WebSocketEvent {
Connected,
Message(String),
Binary(Vec<u8>),
Closed(Option<String>),
Error(WebSocketError),
}
#[derive(Debug, Clone)]
pub enum WebSocketError {
ConnectionFailed(String),
SendFailed(String),
ProtocolError(String),
Timeout,
}
impl std::fmt::Display for WebSocketError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::ConnectionFailed(e) => write!(f, "Connection failed: {e}"),
Self::SendFailed(e) => write!(f, "Send failed: {e}"),
Self::ProtocolError(e) => write!(f, "Protocol error: {e}"),
Self::Timeout => write!(f, "WebSocket timeout"),
}
}
}
impl std::error::Error for WebSocketError {}
#[derive(Debug, Clone)]
pub struct WebSocketHandle {
pub id: String,
pub url: String,
pub connected: Arc<Mutex<bool>>,
}
impl WebSocketHandle {
pub async fn send_text(&self, _message: String) -> Result<(), WebSocketError> {
if *self.connected.lock().await {
Ok(())
} else {
Err(WebSocketError::SendFailed("Not connected".to_string()))
}
}
pub async fn send_binary(&self, _data: Vec<u8>) -> Result<(), WebSocketError> {
if *self.connected.lock().await {
Ok(())
} else {
Err(WebSocketError::SendFailed("Not connected".to_string()))
}
}
pub async fn close(&self) {
*self.connected.lock().await = false;
}
}
pub fn websocket<M, F>(url: impl Into<String>, mut handler: F) -> Cmd<M>
where
M: Message,
F: FnMut(WebSocketEvent) -> Option<M> + Send + 'static,
{
let url = url.into();
let handle = WebSocketHandle {
id: uuid::Uuid::as_string(),
url,
connected: Arc::new(Mutex::new(false)),
};
commands::spawn(async move {
tokio::time::sleep(std::time::Duration::from_millis(100)).await;
*handle.connected.lock().await = true;
handler(WebSocketEvent::Connected)
})
}
pub fn websocket_with_heartbeat<M, F>(
url: impl Into<String>,
_ping_interval: std::time::Duration,
mut handler: F,
) -> Cmd<M>
where
M: Message,
F: FnMut(WebSocketEvent) -> Option<M> + Send + 'static,
{
let _url = url.into();
commands::spawn(async move {
tokio::time::sleep(std::time::Duration::from_millis(100)).await;
handler(WebSocketEvent::Connected)
})
}
#[must_use]
pub fn ws_send<M>(handle: WebSocketHandle, message: String) -> Cmd<M>
where
M: Message,
{
commands::spawn(async move {
let _ = handle.send_text(message).await;
None
})
}
#[must_use]
pub fn ws_close<M>(handle: WebSocketHandle) -> Cmd<M>
where
M: Message,
{
commands::spawn(async move {
handle.close().await;
None
})
}
mod uuid {
pub struct Uuid;
impl Uuid {
#[allow(dead_code)]
pub const fn new_v4() -> Self {
Self
}
pub fn as_string() -> String {
format!("ws-{}", std::process::id())
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use proptest::prelude::*;
#[test]
fn test_websocket_error_display() {
let errors = vec![
WebSocketError::ConnectionFailed("Network unreachable".to_string()),
WebSocketError::SendFailed("Socket closed".to_string()),
WebSocketError::ProtocolError("Invalid frame".to_string()),
WebSocketError::Timeout,
];
for error in errors {
let display = error.to_string();
assert!(!display.is_empty());
match error {
WebSocketError::ConnectionFailed(ref msg) => assert!(display.contains(msg)),
WebSocketError::SendFailed(ref msg) => assert!(display.contains(msg)),
WebSocketError::ProtocolError(ref msg) => assert!(display.contains(msg)),
WebSocketError::Timeout => assert!(display.contains("timeout")),
}
}
}
proptest! {
#[test]
fn prop_websocket_error_consistency(error_msg in ".*") {
let error = WebSocketError::ConnectionFailed(error_msg.clone());
let error_string = error.to_string();
prop_assert!(error_string.contains(&error_msg));
}
}
}