use std::sync::Arc;
use std::time::Duration;
use tokio::sync::Mutex;
use tokio_util::sync::CancellationToken;
use sz_orm_mqtt::{MqttConfig, MqttError, MqttPlugin, QoS};
#[derive(Debug, Clone)]
pub struct MqttRuntimeConfig {
pub client_id: String,
pub keep_alive_secs: u16,
pub topics: Vec<String>,
pub broker_url: String,
}
impl Default for MqttRuntimeConfig {
fn default() -> Self {
Self {
client_id: "sz-rust-mqtt".to_string(),
keep_alive_secs: 60,
topics: Vec::new(),
broker_url: "tcp://localhost:1883".to_string(),
}
}
}
impl MqttRuntimeConfig {
pub fn new(client_id: impl Into<String>) -> Self {
Self {
client_id: client_id.into(),
..Default::default()
}
}
pub fn with_keep_alive(mut self, secs: u16) -> Self {
self.keep_alive_secs = secs;
self
}
pub fn with_topic(mut self, topic: impl Into<String>) -> Self {
self.topics.push(topic.into());
self
}
pub fn with_broker_url(mut self, url: impl Into<String>) -> Self {
self.broker_url = url.into();
self
}
}
pub struct MqttRuntime {
config: MqttRuntimeConfig,
plugin: Arc<Mutex<MqttPlugin>>,
}
impl MqttRuntime {
pub fn new(config: MqttRuntimeConfig) -> Self {
let mqtt_config = MqttConfig::new(config.broker_url.clone())
.with_client_id(config.client_id.clone())
.with_keep_alive(config.keep_alive_secs);
let plugin = MqttPlugin::new(mqtt_config);
Self {
config,
plugin: Arc::new(Mutex::new(plugin)),
}
}
pub async fn connect(&self) -> Result<(), MqttError> {
let mut plugin = self.plugin.lock().await;
plugin.connect().await
}
pub async fn disconnect(&self) -> Result<(), MqttError> {
let mut plugin = self.plugin.lock().await;
plugin.disconnect().await
}
pub async fn is_connected(&self) -> bool {
let plugin = self.plugin.lock().await;
plugin.is_connected()
}
pub async fn subscribe(&self, topic: &str, qos: QoS) -> Result<(), MqttError> {
let plugin = self.plugin.lock().await;
plugin.subscribe(topic, qos).await
}
pub async fn unsubscribe(&self, topic: &str) -> Result<(), MqttError> {
let plugin = self.plugin.lock().await;
plugin.unsubscribe(topic).await
}
pub async fn publish(&self, topic: &str, payload: Vec<u8>, qos: QoS) -> Result<(), MqttError> {
let plugin = self.plugin.lock().await;
plugin.publish(topic, payload, qos).await
}
pub fn start_keepalive(&self, token: CancellationToken) -> tokio::task::JoinHandle<()> {
let plugin = self.plugin.clone();
let interval_secs = self.config.keep_alive_secs.max(1) as u64;
tokio::spawn(async move {
let mut ticker = tokio::time::interval(Duration::from_secs(interval_secs));
loop {
tokio::select! {
_ = token.cancelled() => break,
_ = ticker.tick() => {
let mut p = plugin.lock().await;
if !p.is_connected() {
if let Err(e) = p.connect().await {
tracing::warn!("mqtt reconnect failed: {}", e);
}
}
}
}
}
})
}
pub async fn subscribe_default_topics(&self) -> Result<(), MqttError> {
for topic in &self.config.topics {
self.subscribe(topic, QoS::AtLeastOnce).await?;
}
Ok(())
}
pub fn config(&self) -> &MqttRuntimeConfig {
&self.config
}
pub async fn subscription_count(&self) -> usize {
let plugin = self.plugin.lock().await;
plugin.subscription_count().await
}
pub async fn message_count(&self) -> usize {
let plugin = self.plugin.lock().await;
plugin.message_count().await
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_mqtt_runtime_config_default() {
let config = MqttRuntimeConfig::default();
assert_eq!(config.client_id, "sz-rust-mqtt");
assert_eq!(config.keep_alive_secs, 60);
assert!(config.topics.is_empty());
assert_eq!(config.broker_url, "tcp://localhost:1883");
}
#[test]
fn test_mqtt_runtime_config_builder() {
let config = MqttRuntimeConfig::new("client-1")
.with_keep_alive(30)
.with_topic("orders/#")
.with_topic("payments/#")
.with_broker_url("ssl://broker.example.com:8883");
assert_eq!(config.client_id, "client-1");
assert_eq!(config.keep_alive_secs, 30);
assert_eq!(config.topics.len(), 2);
assert_eq!(config.broker_url, "ssl://broker.example.com:8883");
}
#[tokio::test]
async fn test_mqtt_connect_disconnect() {
let runtime = MqttRuntime::new(MqttRuntimeConfig::new("test-client"));
assert!(!runtime.is_connected().await);
runtime.connect().await.unwrap();
assert!(runtime.is_connected().await);
runtime.disconnect().await.unwrap();
assert!(!runtime.is_connected().await);
}
#[tokio::test]
async fn test_mqtt_subscribe_unsubscribe() {
let runtime = MqttRuntime::new(MqttRuntimeConfig::new("test-client"));
runtime.connect().await.unwrap();
runtime
.subscribe("test/topic", QoS::AtLeastOnce)
.await
.unwrap();
assert_eq!(runtime.subscription_count().await, 1);
runtime.unsubscribe("test/topic").await.unwrap();
assert_eq!(runtime.subscription_count().await, 0);
}
#[tokio::test]
async fn test_mqtt_publish() {
let runtime = MqttRuntime::new(MqttRuntimeConfig::new("test-client"));
runtime.connect().await.unwrap();
runtime
.subscribe("test/topic", QoS::AtLeastOnce)
.await
.unwrap();
runtime
.publish("test/topic", b"hello mqtt".to_vec(), QoS::AtLeastOnce)
.await
.unwrap();
assert_eq!(runtime.message_count().await, 1);
}
#[tokio::test]
async fn test_mqtt_publish_multiple() {
let runtime = MqttRuntime::new(MqttRuntimeConfig::new("test-client"));
runtime.connect().await.unwrap();
runtime
.subscribe("test/topic", QoS::AtLeastOnce)
.await
.unwrap();
for i in 0..5 {
runtime
.publish(
"test/topic",
format!("msg-{}", i).into_bytes(),
QoS::AtLeastOnce,
)
.await
.unwrap();
}
assert_eq!(runtime.message_count().await, 5);
}
#[tokio::test]
async fn test_subscribe_default_topics() {
let config = MqttRuntimeConfig::new("test-client")
.with_topic("orders/#")
.with_topic("payments/#");
let runtime = MqttRuntime::new(config);
runtime.connect().await.unwrap();
runtime.subscribe_default_topics().await.unwrap();
assert_eq!(runtime.subscription_count().await, 2);
}
#[tokio::test]
async fn test_keepalive_task_stops_on_cancel() {
let runtime = MqttRuntime::new(MqttRuntimeConfig::new("test-client").with_keep_alive(1));
let token = CancellationToken::new();
let handle = runtime.start_keepalive(token.clone());
tokio::time::sleep(Duration::from_millis(100)).await;
token.cancel();
let result = tokio::time::timeout(Duration::from_secs(2), handle).await;
assert!(result.is_ok(), "keepalive task should stop on cancel");
}
#[tokio::test]
async fn test_keepalive_reconnects_after_disconnect() {
let runtime = MqttRuntime::new(MqttRuntimeConfig::new("test-client").with_keep_alive(1));
runtime.connect().await.unwrap();
assert!(runtime.is_connected().await);
let token = CancellationToken::new();
let plugin_clone = runtime.plugin.clone();
let handle = runtime.start_keepalive(token.clone());
{
let mut p = plugin_clone.lock().await;
p.disconnect().await.unwrap();
}
assert!(!runtime.is_connected().await);
tokio::time::sleep(Duration::from_millis(1500)).await;
assert!(runtime.is_connected().await);
token.cancel();
let _ = handle.await;
}
#[test]
fn test_config_accessor() {
let runtime = MqttRuntime::new(MqttRuntimeConfig::new("test").with_keep_alive(45));
assert_eq!(runtime.config().client_id, "test");
assert_eq!(runtime.config().keep_alive_secs, 45);
}
#[tokio::test]
async fn test_publish_without_connect_returns_error() {
let runtime = MqttRuntime::new(MqttRuntimeConfig::new("test-client"));
let result = runtime
.publish("test", b"data".to_vec(), QoS::AtMostOnce)
.await;
let _ = result;
}
}