use crate::error::Result;
use serde::de::DeserializeOwned;
use tokio_tungstenite::tungstenite::Message as TungsteniteMessage;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum MessageType {
Text,
Binary,
Ping,
Pong,
Close,
}
#[derive(Debug, Clone)]
pub struct Message {
pub data: Vec<u8>,
pub msg_type: MessageType,
}
impl Message {
pub fn text(text: impl Into<String>) -> Self {
let string = text.into();
Self {
data: string.into_bytes(),
msg_type: MessageType::Text,
}
}
pub fn binary(data: Vec<u8>) -> Self {
Self {
data,
msg_type: MessageType::Binary,
}
}
pub fn ping(data: Vec<u8>) -> Self {
Self {
data,
msg_type: MessageType::Ping,
}
}
pub fn pong(data: Vec<u8>) -> Self {
Self {
data,
msg_type: MessageType::Pong,
}
}
pub fn close() -> Self {
Self {
data: Vec::new(),
msg_type: MessageType::Close,
}
}
pub fn into_tungstenite(self) -> TungsteniteMessage {
match self.msg_type {
MessageType::Text => {
TungsteniteMessage::Text(String::from_utf8_lossy(&self.data).to_string())
}
MessageType::Binary => TungsteniteMessage::Binary(self.data),
MessageType::Ping => TungsteniteMessage::Ping(self.data),
MessageType::Pong => TungsteniteMessage::Pong(self.data),
MessageType::Close => TungsteniteMessage::Close(None),
}
}
pub fn from_tungstenite(msg: TungsteniteMessage) -> Self {
match msg {
TungsteniteMessage::Text(text) => Self::text(text),
TungsteniteMessage::Binary(data) => Self::binary(data),
TungsteniteMessage::Ping(data) => Self::ping(data),
TungsteniteMessage::Pong(data) => Self::pong(data),
TungsteniteMessage::Close(_) => Self::close(),
TungsteniteMessage::Frame(_) => Self::binary(vec![]),
}
}
pub fn message_type(&self) -> MessageType {
self.msg_type
}
pub fn is_text(&self) -> bool {
self.msg_type == MessageType::Text
}
pub fn is_binary(&self) -> bool {
self.msg_type == MessageType::Binary
}
pub fn is_ping(&self) -> bool {
self.msg_type == MessageType::Ping
}
pub fn is_pong(&self) -> bool {
self.msg_type == MessageType::Pong
}
pub fn is_close(&self) -> bool {
self.msg_type == MessageType::Close
}
pub fn as_text(&self) -> Option<&str> {
if self.is_text() {
std::str::from_utf8(&self.data).ok()
} else {
None
}
}
pub fn as_bytes(&self) -> &[u8] {
&self.data
}
pub fn json<T: DeserializeOwned>(&self) -> Result<T> {
let text = self
.as_text()
.ok_or_else(|| crate::error::Error::InvalidMessage)?;
Ok(serde_json::from_str(text)?)
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_text_message() {
let msg = Message::text("Hello, World!");
assert!(msg.is_text());
assert_eq!(msg.as_text(), Some("Hello, World!"));
assert_eq!(msg.message_type(), MessageType::Text);
}
#[test]
fn test_binary_message() {
let data = vec![1, 2, 3, 4, 5];
let msg = Message::binary(data.clone());
assert!(msg.is_binary());
assert_eq!(msg.as_bytes(), &data[..]);
assert_eq!(msg.message_type(), MessageType::Binary);
}
#[test]
fn test_ping_message() {
let msg = Message::ping(vec![]);
assert!(msg.is_ping());
assert_eq!(msg.message_type(), MessageType::Ping);
}
#[test]
fn test_pong_message() {
let msg = Message::pong(vec![]);
assert!(msg.is_pong());
assert_eq!(msg.message_type(), MessageType::Pong);
}
#[test]
fn test_close_message() {
let msg = Message::close();
assert!(msg.is_close());
assert_eq!(msg.message_type(), MessageType::Close);
}
#[test]
fn test_json_parsing() {
let msg = Message::text(r#"{"key":"value","number":42}"#);
let json: serde_json::Value = msg.json().unwrap();
assert_eq!(json["key"], "value");
assert_eq!(json["number"], 42);
}
#[test]
fn test_as_bytes() {
let text_msg = Message::text("Hello");
assert_eq!(text_msg.as_bytes(), b"Hello");
let binary_msg = Message::binary(vec![1, 2, 3]);
assert_eq!(binary_msg.as_bytes(), &[1, 2, 3]);
}
#[test]
fn test_tungstenite_conversion() {
let msg = Message::text("test");
let tung_msg = msg.clone().into_tungstenite();
let back = Message::from_tungstenite(tung_msg);
assert_eq!(back.as_text(), msg.as_text());
}
}