use futures_util::{SinkExt, StreamExt};
use qrusty::{api::ApiServer, storage::Storage};
use reqwest::Client;
use serde_json::{json, Value};
use std::sync::Arc;
use tempfile::TempDir;
use tokio::net::TcpListener;
use tokio::time::{timeout, Duration};
use tokio_tungstenite::{connect_async, tungstenite::Message as TungsteniteMessage};
async fn start_test_server() -> (String, String, TempDir) {
let temp_dir = TempDir::new().unwrap();
let storage = Arc::new(Storage::new(temp_dir.path().to_str().unwrap()).unwrap());
let api = ApiServer::new(storage);
let router = api
.router()
.layer(tower_http::cors::CorsLayer::permissive());
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let port = listener.local_addr().unwrap().port();
tokio::spawn(async move {
axum::serve(listener, router).await.unwrap();
});
(
format!("ws://127.0.0.1:{}", port),
format!("http://127.0.0.1:{}", port),
temp_dir,
)
}
type WsStream =
tokio_tungstenite::WebSocketStream<tokio_tungstenite::MaybeTlsStream<tokio::net::TcpStream>>;
async fn ws_connect(ws_base: &str) -> WsStream {
let url = format!("{}/ws", ws_base);
let (ws, _) = connect_async(&url).await.unwrap();
ws
}
async fn send(ws: &mut WsStream, frame: Value) {
ws.send(TungsteniteMessage::Text(frame.to_string().into()))
.await
.unwrap();
}
async fn recv(ws: &mut WsStream) -> Value {
timeout(Duration::from_secs(5), async {
loop {
let msg = ws.next().await.expect("stream ended").expect("ws error");
match msg {
TungsteniteMessage::Text(t) => {
return serde_json::from_str(&t).expect("bad JSON");
}
TungsteniteMessage::Ping(_) | TungsteniteMessage::Pong(_) => continue,
other => panic!("expected text frame, got {:?}", other),
}
}
})
.await
.expect("recv timeout")
}
async fn http_create_queue(http_base: &str, name: &str) {
Client::new()
.post(format!("{}/create-queue", http_base))
.json(&json!({"name": name, "config": {"ordering": "MaxFirst"}}))
.send()
.await
.unwrap()
.error_for_status()
.unwrap();
}
async fn http_publish(http_base: &str, queue: &str, payload: &str) -> String {
let resp = Client::new()
.post(format!("{}/publish", http_base))
.json(&json!({"queue": queue, "priority": 0, "payload": payload}))
.send()
.await
.unwrap()
.error_for_status()
.unwrap();
let body: Value = resp.json().await.unwrap();
body["id"].as_str().unwrap().to_string()
}
#[tokio::test]
async fn test_ws_connection_upgrade() {
let (ws_base, _http, _dir) = start_test_server().await;
let mut ws = ws_connect(&ws_base).await;
ws.close(None).await.unwrap();
}
#[tokio::test]
async fn test_ws_binary_frame_rejected_with_1003() {
let (ws_base, _http, _dir) = start_test_server().await;
let mut ws = ws_connect(&ws_base).await;
ws.send(TungsteniteMessage::Binary(b"raw bytes".to_vec().into()))
.await
.unwrap();
let msg = timeout(Duration::from_secs(2), ws.next())
.await
.expect("timeout waiting for close")
.expect("stream ended")
.expect("ws error");
match msg {
TungsteniteMessage::Close(Some(frame)) => {
assert_eq!(
frame.code,
tokio_tungstenite::tungstenite::protocol::frame::coding::CloseCode::Unsupported
);
}
other => panic!("expected close frame, got {:?}", other),
}
}
#[tokio::test]
async fn test_ws_invalid_json_rejected_with_1007() {
let (ws_base, _http, _dir) = start_test_server().await;
let mut ws = ws_connect(&ws_base).await;
ws.send(TungsteniteMessage::Text("{not valid json".into()))
.await
.unwrap();
let msg = timeout(Duration::from_secs(2), ws.next())
.await
.expect("timeout waiting for close")
.expect("stream ended")
.expect("ws error");
match msg {
TungsteniteMessage::Close(Some(frame)) => {
assert_eq!(
frame.code,
tokio_tungstenite::tungstenite::protocol::frame::coding::CloseCode::Invalid
);
}
other => panic!("expected close frame, got {:?}", other),
}
}
#[tokio::test]
async fn test_ws_req_id_echoed_in_ok() {
let (ws_base, http_base, _dir) = start_test_server().await;
http_create_queue(&http_base, "echo-q").await;
let mut ws = ws_connect(&ws_base).await;
send(
&mut ws,
json!({
"type": "publish",
"queue": "echo-q",
"payload": "hello",
"req_id": "my-request-42"
}),
)
.await;
let resp = recv(&mut ws).await;
assert_eq!(resp["type"], "ok");
assert_eq!(resp["req_id"], "my-request-42");
assert!(resp["id"].is_string());
}
#[tokio::test]
async fn test_ws_no_req_id_omitted_in_response() {
let (ws_base, http_base, _dir) = start_test_server().await;
http_create_queue(&http_base, "no-rid-q").await;
let mut ws = ws_connect(&ws_base).await;
send(
&mut ws,
json!({"type": "publish", "queue": "no-rid-q", "payload": "x"}),
)
.await;
let resp = recv(&mut ws).await;
assert_eq!(resp["type"], "ok");
assert!(resp.get("req_id").is_none(), "req_id should be absent");
}
#[tokio::test]
async fn test_ws_unknown_type_returns_error_frame_and_stays_open() {
let (ws_base, _http, _dir) = start_test_server().await;
let mut ws = ws_connect(&ws_base).await;
send(&mut ws, json!({"type": "totally-unknown", "req_id": "r1"})).await;
let resp = recv(&mut ws).await;
assert_eq!(resp["type"], "error");
assert_eq!(resp["req_id"], "r1");
assert!(resp["code"].is_string());
assert!(resp["message"].is_string());
send(&mut ws, json!({"type": "totally-unknown"})).await;
let resp2 = recv(&mut ws).await;
assert_eq!(resp2["type"], "error");
}
#[tokio::test]
async fn test_ws_publish_success() {
let (ws_base, http_base, _dir) = start_test_server().await;
http_create_queue(&http_base, "pub-q").await;
let mut ws = ws_connect(&ws_base).await;
send(
&mut ws,
json!({"type": "publish", "queue": "pub-q", "payload": "world", "priority": 5}),
)
.await;
let resp = recv(&mut ws).await;
assert_eq!(resp["type"], "ok");
assert!(resp["id"].is_string(), "response should include message id");
}
#[tokio::test]
async fn test_ws_publish_nonexistent_queue_error() {
let (ws_base, _http, _dir) = start_test_server().await;
let mut ws = ws_connect(&ws_base).await;
send(
&mut ws,
json!({"type": "publish", "queue": "ghost-q", "payload": "x", "req_id": "r99"}),
)
.await;
let resp = recv(&mut ws).await;
assert_eq!(resp["type"], "error");
assert_eq!(resp["req_id"], "r99");
assert!(resp["code"].is_string());
send(&mut ws, json!({"type": "totally-unknown"})).await;
let still_open = recv(&mut ws).await;
assert_eq!(still_open["type"], "error");
}
#[tokio::test]
async fn test_ws_subscribe_and_deliver() {
let (ws_base, http_base, _dir) = start_test_server().await;
http_create_queue(&http_base, "sub-q").await;
let mut ws = ws_connect(&ws_base).await;
send(
&mut ws,
json!({"type": "subscribe", "queue": "sub-q", "req_id": "sub1"}),
)
.await;
let sub_resp = recv(&mut ws).await;
assert_eq!(sub_resp["type"], "ok");
assert_eq!(sub_resp["req_id"], "sub1");
http_publish(&http_base, "sub-q", "hello-ws").await;
let delivery = timeout(Duration::from_secs(3), async {
loop {
let frame = recv(&mut ws).await;
if frame["type"] == "deliver" {
return frame;
}
}
})
.await
.expect("did not receive delivery frame in time");
assert_eq!(delivery["type"], "deliver");
assert_eq!(delivery["queue"], "sub-q");
assert_eq!(delivery["payload"], "hello-ws");
assert!(delivery["id"].is_string());
assert!(delivery["priority"].is_number());
assert!(delivery["created_at"].is_string());
assert!(
delivery.get("req_id").is_none(),
"delivery frame must not include req_id"
);
}
#[tokio::test]
async fn test_ws_subscribe_nonexistent_queue_error() {
let (ws_base, _http, _dir) = start_test_server().await;
let mut ws = ws_connect(&ws_base).await;
send(
&mut ws,
json!({"type": "subscribe", "queue": "no-such-q", "req_id": "s1"}),
)
.await;
let resp = recv(&mut ws).await;
assert_eq!(resp["type"], "error");
assert_eq!(resp["req_id"], "s1");
}
#[tokio::test]
async fn test_ws_unsubscribe_stops_delivery() {
let (ws_base, http_base, _dir) = start_test_server().await;
http_create_queue(&http_base, "unsub-q").await;
let mut ws = ws_connect(&ws_base).await;
send(&mut ws, json!({"type": "subscribe", "queue": "unsub-q"})).await;
let _ = recv(&mut ws).await;
send(
&mut ws,
json!({"type": "unsubscribe", "queue": "unsub-q", "req_id": "u1"}),
)
.await;
let unsub_resp = recv(&mut ws).await;
assert_eq!(unsub_resp["type"], "ok");
assert_eq!(unsub_resp["req_id"], "u1");
http_publish(&http_base, "unsub-q", "should-not-arrive").await;
let no_delivery = timeout(Duration::from_millis(400), ws.next()).await;
assert!(
no_delivery.is_err(),
"should not have received a frame after unsubscribe"
);
}
#[tokio::test]
async fn test_ws_multi_queue_subscribe() {
let (ws_base, http_base, _dir) = start_test_server().await;
http_create_queue(&http_base, "mq-alpha").await;
http_create_queue(&http_base, "mq-beta").await;
let mut ws = ws_connect(&ws_base).await;
send(&mut ws, json!({"type": "subscribe", "queue": "mq-alpha"})).await;
let _ = recv(&mut ws).await; send(&mut ws, json!({"type": "subscribe", "queue": "mq-beta"})).await;
let _ = recv(&mut ws).await;
http_publish(&http_base, "mq-alpha", "from-alpha").await;
http_publish(&http_base, "mq-beta", "from-beta").await;
let mut queues_seen = std::collections::HashSet::new();
for _ in 0..2 {
let frame = timeout(Duration::from_secs(3), async {
loop {
let f = recv(&mut ws).await;
if f["type"] == "deliver" {
return f;
}
}
})
.await
.expect("timeout waiting for delivery");
queues_seen.insert(frame["queue"].as_str().unwrap().to_string());
}
assert!(
queues_seen.contains("mq-alpha"),
"missing mq-alpha delivery"
);
assert!(queues_seen.contains("mq-beta"), "missing mq-beta delivery");
}
#[tokio::test]
async fn test_ws_ack_removes_message() {
let (ws_base, http_base, _dir) = start_test_server().await;
http_create_queue(&http_base, "ack-q").await;
let mut ws = ws_connect(&ws_base).await;
send(&mut ws, json!({"type": "subscribe", "queue": "ack-q"})).await;
let _ = recv(&mut ws).await;
http_publish(&http_base, "ack-q", "ack-me").await;
let delivery = timeout(Duration::from_secs(3), async {
loop {
let f = recv(&mut ws).await;
if f["type"] == "deliver" {
return f;
}
}
})
.await
.unwrap();
let msg_id = delivery["id"].as_str().unwrap().to_string();
send(
&mut ws,
json!({"type": "ack", "queue": "ack-q", "id": msg_id, "req_id": "a1"}),
)
.await;
let ack_resp = recv(&mut ws).await;
assert_eq!(ack_resp["type"], "ok");
assert_eq!(ack_resp["req_id"], "a1");
let stats: Value = Client::new()
.get(format!("{}/queue-stats/ack-q", http_base))
.send()
.await
.unwrap()
.json()
.await
.unwrap();
assert_eq!(stats["total"], 0, "message should be gone after ack");
}
#[tokio::test]
async fn test_ws_nack_unlocks_message() {
let (ws_base, http_base, _dir) = start_test_server().await;
http_create_queue(&http_base, "nack-q").await;
let mut ws = ws_connect(&ws_base).await;
send(&mut ws, json!({"type": "subscribe", "queue": "nack-q"})).await;
let _ = recv(&mut ws).await;
http_publish(&http_base, "nack-q", "nack-me").await;
let delivery = timeout(Duration::from_secs(3), async {
loop {
let f = recv(&mut ws).await;
if f["type"] == "deliver" {
return f;
}
}
})
.await
.unwrap();
let msg_id = delivery["id"].as_str().unwrap().to_string();
send(
&mut ws,
json!({"type": "nack", "queue": "nack-q", "id": msg_id, "req_id": "n1"}),
)
.await;
let nack_resp = loop {
let f = recv(&mut ws).await;
if f["type"] == "ok" {
break f;
}
assert_eq!(f["type"], "deliver", "unexpected frame: {f}");
};
assert_eq!(nack_resp["req_id"], "n1");
send(
&mut ws,
json!({"type": "unsubscribe", "queue": "nack-q", "req_id": "u1"}),
)
.await;
loop {
let f = recv(&mut ws).await;
if f["type"] == "ok" && f["req_id"] == "u1" {
break;
}
}
let stats: Value = Client::new()
.get(format!("{}/queue-stats/nack-q", http_base))
.send()
.await
.unwrap()
.json()
.await
.unwrap();
assert!(
stats["available"].as_u64().unwrap_or(0) > 0,
"message should be back in the queue after nack"
);
}
#[tokio::test]
async fn test_ws_batch_ack() {
let (ws_base, http_base, _dir) = start_test_server().await;
http_create_queue(&http_base, "batch-ack-q").await;
let mut ws = ws_connect(&ws_base).await;
http_publish(&http_base, "batch-ack-q", "m1").await;
http_publish(&http_base, "batch-ack-q", "m2").await;
send(
&mut ws,
json!({"type": "subscribe", "queue": "batch-ack-q"}),
)
.await;
let _ = recv(&mut ws).await;
let mut ids = vec![];
for _ in 0..2 {
let f = timeout(Duration::from_secs(3), async {
loop {
let f = recv(&mut ws).await;
if f["type"] == "deliver" {
return f;
}
}
})
.await
.unwrap();
ids.push(f["id"].as_str().unwrap().to_string());
}
send(
&mut ws,
json!({"type": "batch-ack", "queue": "batch-ack-q", "ids": ids, "req_id": "ba1"}),
)
.await;
let resp = recv(&mut ws).await;
assert_eq!(resp["type"], "ok");
assert_eq!(resp["req_id"], "ba1");
assert_eq!(resp["acked"], 2, "both messages should be acked");
}
#[tokio::test]
async fn test_ws_batch_nack() {
let (ws_base, http_base, _dir) = start_test_server().await;
http_create_queue(&http_base, "batch-nack-q").await;
let mut ws = ws_connect(&ws_base).await;
http_publish(&http_base, "batch-nack-q", "n1").await;
http_publish(&http_base, "batch-nack-q", "n2").await;
send(
&mut ws,
json!({"type": "subscribe", "queue": "batch-nack-q"}),
)
.await;
let _ = recv(&mut ws).await;
let mut ids = vec![];
for _ in 0..2 {
let f = timeout(Duration::from_secs(3), async {
loop {
let f = recv(&mut ws).await;
if f["type"] == "deliver" {
return f;
}
}
})
.await
.unwrap();
ids.push(f["id"].as_str().unwrap().to_string());
}
send(
&mut ws,
json!({"type": "batch-nack", "queue": "batch-nack-q", "ids": ids, "req_id": "bn1"}),
)
.await;
let resp = recv(&mut ws).await;
assert_eq!(resp["type"], "ok");
assert_eq!(resp["req_id"], "bn1");
let unlocked = resp["unlocked"].as_u64().unwrap_or(0);
let dropped = resp["dropped"].as_u64().unwrap_or(0);
assert_eq!(
unlocked + dropped,
2,
"all messages should be accounted for"
);
}
#[tokio::test]
async fn test_ws_connection_close_releases_subscriptions() {
let (ws_base, http_base, _dir) = start_test_server().await;
http_create_queue(&http_base, "close-q").await;
let mut ws = ws_connect(&ws_base).await;
send(&mut ws, json!({"type": "subscribe", "queue": "close-q"})).await;
let _ = recv(&mut ws).await;
ws.close(None).await.unwrap();
tokio::time::sleep(Duration::from_millis(200)).await;
let status = Client::new()
.get(format!("{}/health", http_base))
.send()
.await
.unwrap()
.status();
assert_eq!(status, 200);
}
#[tokio::test]
async fn test_ws_error_frame_echoes_req_id() {
let (ws_base, _http, _dir) = start_test_server().await;
let mut ws = ws_connect(&ws_base).await;
send(
&mut ws,
json!({"type": "subscribe", "queue": "nonexistent-q", "req_id": "err-123"}),
)
.await;
let resp = recv(&mut ws).await;
assert_eq!(resp["type"], "error");
assert_eq!(resp["req_id"], "err-123");
assert!(resp["code"].is_string());
assert!(resp["message"].is_string());
}
#[tokio::test]
async fn test_ws_publish_missing_queue_field() {
let (ws_base, _http, _dir) = start_test_server().await;
let mut ws = ws_connect(&ws_base).await;
send(&mut ws, json!({"type": "publish", "payload": "oops"})).await;
let resp = recv(&mut ws).await;
assert_eq!(resp["type"], "error");
}
#[tokio::test]
async fn test_ws_renew_extends_lock() {
let (ws_base, http_base, _dir) = start_test_server().await;
http_create_queue(&http_base, "renew-q").await;
let mut ws = ws_connect(&ws_base).await;
send(&mut ws, json!({"type": "subscribe", "queue": "renew-q"})).await;
let _ = recv(&mut ws).await;
http_publish(&http_base, "renew-q", "renew-me").await;
let delivery = timeout(Duration::from_secs(3), async {
loop {
let f = recv(&mut ws).await;
if f["type"] == "deliver" {
return f;
}
}
})
.await
.expect("did not receive delivery frame in time");
let msg_id = delivery["id"].as_str().unwrap().to_string();
send(
&mut ws,
json!({"type": "renew", "queue": "renew-q", "id": msg_id, "req_id": "rn1"}),
)
.await;
let renew_resp = recv(&mut ws).await;
assert_eq!(renew_resp["type"], "ok");
assert_eq!(renew_resp["req_id"], "rn1");
send(
&mut ws,
json!({"type": "ack", "queue": "renew-q", "id": msg_id}),
)
.await;
let ack_resp = recv(&mut ws).await;
assert_eq!(ack_resp["type"], "ok");
}
#[tokio::test]
async fn test_ws_renew_not_locked_returns_error() {
let (ws_base, http_base, _dir) = start_test_server().await;
http_create_queue(&http_base, "renew-err-q").await;
let mut ws = ws_connect(&ws_base).await;
let msg_id = http_publish(&http_base, "renew-err-q", "not-locked").await;
send(
&mut ws,
json!({"type": "renew", "queue": "renew-err-q", "id": msg_id, "req_id": "rn2"}),
)
.await;
let resp = recv(&mut ws).await;
assert_eq!(resp["type"], "error");
assert_eq!(resp["code"], "not_locked");
assert_eq!(resp["req_id"], "rn2");
send(&mut ws, json!({"type": "totally-unknown"})).await;
let still_open = recv(&mut ws).await;
assert_eq!(still_open["type"], "error");
}
#[tokio::test]
async fn test_ws_operations_update_rate_counters() {
let (ws_base, http_base, _dir) = start_test_server().await;
http_create_queue(&http_base, "rate-ws").await;
let _: Value = Client::new()
.get(format!("{}/queue-stats/rate-ws", http_base))
.send()
.await
.unwrap()
.json()
.await
.unwrap();
let mut ws = ws_connect(&ws_base).await;
for i in 0..3u64 {
send(
&mut ws,
json!({"type": "publish", "queue": "rate-ws", "priority": i, "payload": format!("msg-{}", i)}),
)
.await;
let resp = recv(&mut ws).await;
assert_eq!(resp["type"], "ok", "publish frame {i} failed: {resp}");
}
send(&mut ws, json!({"type": "subscribe", "queue": "rate-ws"})).await;
let sub_resp = recv(&mut ws).await;
assert_eq!(sub_resp["type"], "ok");
let mut pending_acks = 0u32;
let mut acked = 0u32;
let total = 3u32;
while acked < total {
let f = timeout(Duration::from_secs(5), recv(&mut ws))
.await
.expect("timed out waiting for deliver/ok frame");
match f["type"].as_str().unwrap_or("") {
"deliver" => {
let msg_id = f["id"].as_str().unwrap().to_owned();
send(
&mut ws,
json!({"type": "ack", "queue": "rate-ws", "id": msg_id}),
)
.await;
pending_acks += 1;
}
"ok" => {
assert!(pending_acks > 0, "unexpected ok with no pending acks");
pending_acks -= 1;
acked += 1;
}
other => panic!("unexpected frame type: {other} — {f}"),
}
}
let deadline = std::time::Instant::now() + Duration::from_secs(3);
let mut pub_rate = 0.0f64;
let mut ack_rate = 0.0f64;
while std::time::Instant::now() < deadline {
let stats: Value = Client::new()
.get(format!("{}/queue-stats/rate-ws", http_base))
.send()
.await
.unwrap()
.json()
.await
.unwrap();
pub_rate = stats["publish_rate_per_sec"].as_f64().unwrap_or(0.0);
ack_rate = stats["ack_rate_per_sec"].as_f64().unwrap_or(0.0);
if pub_rate > 0.0 && ack_rate > 0.0 {
break;
}
tokio::time::sleep(Duration::from_millis(100)).await;
}
assert!(
pub_rate > 0.0,
"publish_rate_per_sec must be non-zero after WS publishes; got {pub_rate}"
);
assert!(
ack_rate > 0.0,
"ack_rate_per_sec must be non-zero after WS acks; got {ack_rate}"
);
}
#[tokio::test]
async fn test_ws_no_reply_suppresses_response() {
let (ws_base, http_base, _dir) = start_test_server().await;
http_create_queue(&http_base, "noreply-q").await;
let mut ws = ws_connect(&ws_base).await;
http_publish(&http_base, "noreply-q", "msg").await;
send(&mut ws, json!({"type": "subscribe", "queue": "noreply-q"})).await;
let _ = recv(&mut ws).await;
let delivery = timeout(Duration::from_secs(3), async {
loop {
let f = recv(&mut ws).await;
if f["type"] == "deliver" {
return f;
}
}
})
.await
.unwrap();
let msg_id = delivery["id"].as_str().unwrap().to_string();
send(
&mut ws,
json!({"type": "unsubscribe", "queue": "noreply-q"}),
)
.await;
let _ = recv(&mut ws).await;
send(
&mut ws,
json!({
"type": "ack",
"queue": "noreply-q",
"id": msg_id,
"no_reply": true
}),
)
.await;
send(
&mut ws,
json!({"type": "totally-unknown", "req_id": "probe"}),
)
.await;
let resp = recv(&mut ws).await;
assert_eq!(
resp["req_id"], "probe",
"no_reply ack should not produce a response; got: {resp}"
);
let stats: Value = Client::new()
.get(format!("{}/queue-stats/noreply-q", http_base))
.send()
.await
.unwrap()
.json()
.await
.unwrap();
assert_eq!(
stats["total"], 0,
"message should be acked despite no_reply"
);
}
#[tokio::test]
async fn test_ws_credit_flow_control() {
let (ws_base, http_base, _dir) = start_test_server().await;
http_create_queue(&http_base, "credit-q").await;
let mut ws = ws_connect(&ws_base).await;
for i in 0..5 {
http_publish(&http_base, "credit-q", &format!("msg-{}", i)).await;
}
send(
&mut ws,
json!({"type": "subscribe", "queue": "credit-q", "credits": 2, "req_id": "sub1"}),
)
.await;
let resp = recv(&mut ws).await;
assert_eq!(resp["type"], "ok");
assert_eq!(resp["req_id"], "sub1");
for _ in 0..2 {
let f = timeout(Duration::from_secs(3), async {
loop {
let f = recv(&mut ws).await;
if f["type"] == "deliver" {
return f;
}
}
})
.await
.expect("expected delivery");
assert_eq!(f["queue"], "credit-q");
}
let no_more = timeout(Duration::from_millis(500), async {
loop {
let f = recv(&mut ws).await;
if f["type"] == "deliver" {
return f;
}
}
})
.await;
assert!(
no_more.is_err(),
"should not receive more than 2 deliveries"
);
send(
&mut ws,
json!({"type": "credit", "queue": "credit-q", "credits": 2, "req_id": "cr1"}),
)
.await;
let cr_resp = recv(&mut ws).await;
assert_eq!(cr_resp["type"], "ok");
assert_eq!(cr_resp["req_id"], "cr1");
for _ in 0..2 {
let f = timeout(Duration::from_secs(3), async {
loop {
let f = recv(&mut ws).await;
if f["type"] == "deliver" {
return f;
}
}
})
.await
.expect("expected delivery after credit replenishment");
assert_eq!(f["queue"], "credit-q");
}
}
#[tokio::test]
async fn test_ws_subscribe_without_credits_is_unlimited() {
let (ws_base, http_base, _dir) = start_test_server().await;
http_create_queue(&http_base, "unlim-q").await;
let mut ws = ws_connect(&ws_base).await;
for i in 0..5 {
http_publish(&http_base, "unlim-q", &format!("msg-{}", i)).await;
}
send(&mut ws, json!({"type": "subscribe", "queue": "unlim-q"})).await;
let _ = recv(&mut ws).await;
for _ in 0..5 {
let f = timeout(Duration::from_secs(5), async {
loop {
let f = recv(&mut ws).await;
if f["type"] == "deliver" {
return f;
}
}
})
.await
.expect("expected delivery");
assert_eq!(f["queue"], "unlim-q");
}
}
#[tokio::test]
async fn test_ws_credit_on_unsubscribed_queue_returns_error() {
let (ws_base, _http_base, _dir) = start_test_server().await;
let mut ws = ws_connect(&ws_base).await;
send(
&mut ws,
json!({"type": "credit", "queue": "no-such-sub", "credits": 5, "req_id": "c1"}),
)
.await;
let resp = recv(&mut ws).await;
assert_eq!(resp["type"], "error");
assert_eq!(resp["code"], "not_subscribed");
}
#[tokio::test]
async fn test_ws_session_send_frame_public() {
let (ws_base, http_base, _dir) = start_test_server().await;
http_create_queue(&http_base, "sf-q").await;
let session = qrusty_client::ws::WsSession::connect(&format!(
"ws://127.0.0.1:{}",
extract_port(&ws_base)
))
.await
.unwrap();
session
.send_frame(json!({
"type": "publish",
"queue": "sf-q",
"payload": "raw-frame",
"priority": 0,
}))
.await
.unwrap();
tokio::time::sleep(Duration::from_millis(200)).await;
let stats: Value = Client::new()
.get(format!("{}/queue-stats/sf-q", http_base))
.send()
.await
.unwrap()
.json()
.await
.unwrap();
assert_eq!(
stats["total"], 1,
"send_frame should have published a message"
);
session.close().await.unwrap();
}
#[tokio::test]
async fn test_ws_session_request_timeout() {
use tokio_tungstenite::accept_async;
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let port = listener.local_addr().unwrap().port();
tokio::spawn(async move {
if let Ok((stream, _)) = listener.accept().await {
let _ws = accept_async(stream).await;
tokio::time::sleep(Duration::from_secs(600)).await;
}
});
let session = qrusty_client::ws::WsSession::connect_with_timeout(
&format!("ws://127.0.0.1:{}", port),
Duration::from_millis(100),
)
.await
.unwrap();
let result = session.publish("any-q", "hello", None).await;
assert!(
matches!(
result,
Err(qrusty_client::error::QrustyClientError::RequestTimeout(_))
),
"expected RequestTimeout, got: {:?}",
result,
);
}
#[tokio::test]
async fn test_ws_session_ack_noreply() {
let (ws_base, http_base, _dir) = start_test_server().await;
http_create_queue(&http_base, "noreply-client-q").await;
let session = qrusty_client::ws::WsSession::connect(&format!(
"ws://127.0.0.1:{}",
extract_port(&ws_base)
))
.await
.unwrap();
session
.publish("noreply-client-q", "msg", None)
.await
.unwrap();
let mut rx = session.subscribe("noreply-client-q").await.unwrap();
let msg = tokio::time::timeout(Duration::from_secs(3), rx.recv())
.await
.expect("delivery timeout")
.expect("channel closed")
.expect("delivery error");
session.ack_noreply(&msg.queue, &msg.id).await.unwrap();
tokio::time::sleep(Duration::from_millis(200)).await;
let stats: Value = Client::new()
.get(format!("{}/queue-stats/noreply-client-q", http_base))
.send()
.await
.unwrap()
.json()
.await
.unwrap();
assert_eq!(stats["total"], 0, "message should be acked via ack_noreply");
}
#[tokio::test]
async fn test_ws_session_full_api_round_trip() {
let (ws_base, http_base, _dir) = start_test_server().await;
http_create_queue(&http_base, "wsapi-q").await;
let session = qrusty_client::ws::WsSession::connect(&format!(
"ws://127.0.0.1:{}",
extract_port(&ws_base)
))
.await
.unwrap();
for i in 0..4 {
let id = session
.publish("wsapi-q", &format!("payload-{i}"), None)
.await
.unwrap();
assert!(!id.is_empty(), "publish must return a non-empty message id");
}
let mut rx = session.subscribe("wsapi-q").await.unwrap();
let mut delivered = Vec::new();
for _ in 0..4 {
let msg = tokio::time::timeout(Duration::from_secs(5), rx.recv())
.await
.expect("delivery timeout")
.expect("subscription channel closed")
.expect("delivery error");
delivered.push(msg);
}
let first = &delivered[0];
assert_eq!(first.queue, "wsapi-q");
assert!(!first.id.is_empty());
assert!(first.payload.starts_with("payload-"));
let _ = &first.priority;
let _ = &first.created_at;
session
.ack(&delivered[0].queue, &delivered[0].id)
.await
.unwrap();
session
.nack(&delivered[1].queue, &delivered[1].id)
.await
.unwrap();
let acked = session
.batch_ack("wsapi-q", &[delivered[2].id.as_str()])
.await
.unwrap();
assert_eq!(acked, 1, "batch_ack should ack the one locked message");
let (unlocked, _dropped) = session
.batch_nack("wsapi-q", &[delivered[3].id.as_str()])
.await
.unwrap();
assert_eq!(
unlocked, 1,
"batch_nack should unlock the one locked message"
);
session.unsubscribe("wsapi-q").await.unwrap();
session.close().await.unwrap();
}
#[tokio::test]
async fn test_ws_session_credit_flow_control() {
let (ws_base, http_base, _dir) = start_test_server().await;
http_create_queue(&http_base, "credit-session-q").await;
let session = qrusty_client::ws::WsSession::connect(&format!(
"ws://127.0.0.1:{}",
extract_port(&ws_base)
))
.await
.unwrap();
for i in 0..4 {
session
.publish("credit-session-q", &format!("msg-{}", i), None)
.await
.unwrap();
}
let mut rx = session
.subscribe_with_credits("credit-session-q", Some(2))
.await
.unwrap();
for _ in 0..2 {
let msg = tokio::time::timeout(Duration::from_secs(3), rx.recv())
.await
.expect("delivery timeout")
.expect("channel closed")
.expect("delivery error");
session.ack_noreply(&msg.queue, &msg.id).await.unwrap();
}
let no_more = tokio::time::timeout(Duration::from_millis(500), rx.recv()).await;
assert!(
no_more.is_err(),
"should not receive more than 2 deliveries"
);
session.grant_credits("credit-session-q", 2).await.unwrap();
for _ in 0..2 {
let msg = tokio::time::timeout(Duration::from_secs(3), rx.recv())
.await
.expect("delivery timeout after credit grant")
.expect("channel closed")
.expect("delivery error");
session.ack_noreply(&msg.queue, &msg.id).await.unwrap();
}
}
async fn raw_ws_upgrade(addr: &str) -> tokio::net::TcpStream {
use tokio::io::{AsyncReadExt, AsyncWriteExt};
let mut stream = tokio::net::TcpStream::connect(addr).await.unwrap();
let req = format!(
"GET /ws HTTP/1.1\r\nHost: {addr}\r\nUpgrade: websocket\r\n\
Connection: Upgrade\r\nSec-WebSocket-Key: dGhlIHNhbXBsZSBub25jZQ==\r\n\
Sec-WebSocket-Version: 13\r\n\r\n"
);
stream.write_all(req.as_bytes()).await.unwrap();
let mut buf = Vec::new();
let mut b = [0u8; 1];
loop {
stream.read_exact(&mut b).await.unwrap();
buf.push(b[0]);
if buf.ends_with(b"\r\n\r\n") {
break;
}
}
let headers = String::from_utf8_lossy(&buf);
assert!(
headers.contains("101"),
"expected 101 Switching Protocols, got: {headers}"
);
stream
}
async fn read_raw_ws_frame(stream: &mut tokio::net::TcpStream) -> (u8, Vec<u8>) {
use tokio::io::AsyncReadExt;
let mut hdr = [0u8; 2];
stream.read_exact(&mut hdr).await.unwrap();
let opcode = hdr[0] & 0x0F;
let len_byte = (hdr[1] & 0x7F) as usize;
let payload_len = match len_byte {
126 => {
let mut e = [0u8; 2];
stream.read_exact(&mut e).await.unwrap();
u16::from_be_bytes(e) as usize
}
127 => {
let mut e = [0u8; 8];
stream.read_exact(&mut e).await.unwrap();
u64::from_be_bytes(e) as usize
}
n => n,
};
let mut payload = vec![0u8; payload_len];
if payload_len > 0 {
stream.read_exact(&mut payload).await.unwrap();
}
(opcode, payload)
}
#[tokio::test]
async fn test_ws_ping_timeout_closes_with_1001() {
#[allow(deprecated)]
unsafe {
std::env::set_var("WS_PING_INTERVAL_SECS", "1");
std::env::set_var("WS_PING_TIMEOUT_SECS", "1");
}
let (ws_base, _http_base, _dir) = start_test_server().await;
let addr = ws_base.trim_start_matches("ws://");
let mut stream = raw_ws_upgrade(addr).await;
let close_code = tokio::time::timeout(Duration::from_secs(8), async {
loop {
let (opcode, payload) = read_raw_ws_frame(&mut stream).await;
match opcode {
0x8 => {
return if payload.len() >= 2 {
u16::from_be_bytes([payload[0], payload[1]])
} else {
0u16
};
}
0x9 => {} _ => {} }
}
})
.await
.expect("timed out: server did not send Close within 8 s");
assert_eq!(close_code, 1001, "expected close code 1001 (Going Away)");
#[allow(deprecated)]
unsafe {
std::env::remove_var("WS_PING_INTERVAL_SECS");
std::env::remove_var("WS_PING_TIMEOUT_SECS");
}
}
fn extract_port(base: &str) -> u16 {
base.rsplit(':').next().unwrap().parse().unwrap()
}
async fn start_test_server_with_logs() -> (String, String, qrusty::log_buffer::LogBuffer, TempDir) {
let temp_dir = TempDir::new().unwrap();
let storage = Arc::new(Storage::new(temp_dir.path().to_str().unwrap()).unwrap());
let log_buffer = qrusty::log_buffer::LogBuffer::new();
let api = ApiServer::new(storage).with_log_buffer(log_buffer.clone());
let router = api
.router()
.layer(tower_http::cors::CorsLayer::permissive());
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
let port = listener.local_addr().unwrap().port();
tokio::spawn(async move {
axum::serve(listener, router).await.unwrap();
});
(
format!("ws://127.0.0.1:{}", port),
format!("http://127.0.0.1:{}", port),
log_buffer,
temp_dir,
)
}
#[tokio::test]
async fn test_ws_subscribe_logs_delivers_live_entries() {
let (ws_base, _http, log_buffer, _dir) = start_test_server_with_logs().await;
let mut ws = ws_connect(&ws_base).await;
log_buffer.push_sync("INFO", "history-entry");
send(&mut ws, json!({"type": "subscribe-logs", "req_id": "sl1"})).await;
let resp = recv(&mut ws).await;
assert_eq!(
resp["type"], "ok",
"subscribe-logs should return ok: {resp}"
);
let log_frame = timeout(Duration::from_secs(2), async {
loop {
let f = recv(&mut ws).await;
if f["type"] == "log" {
return f;
}
}
})
.await
.expect("expected log frame from history");
assert_eq!(log_frame["type"], "log");
assert!(log_frame["message"].is_string());
assert!(log_frame["timestamp"].is_string());
assert!(log_frame["level"].is_string());
log_buffer.push_sync("WARN", "live-entry");
let live_frame = timeout(Duration::from_secs(2), async {
loop {
let f = recv(&mut ws).await;
if f["type"] == "log" && f["message"] == "live-entry" {
return f;
}
}
})
.await
.expect("expected live log frame");
assert_eq!(live_frame["level"], "WARN");
}
#[tokio::test]
async fn test_ws_subscribe_logs_with_large_history_does_not_deadlock() {
let (ws_base, _http, log_buffer, _dir) = start_test_server_with_logs().await;
let mut ws = ws_connect(&ws_base).await;
for i in 0..500 {
log_buffer.push_sync("INFO", format!("hist-{i}"));
}
send(
&mut ws,
json!({"type": "subscribe-logs", "req_id": "sl-big"}),
)
.await;
let resp = recv(&mut ws).await;
assert_eq!(resp["type"], "ok");
let first = b"hist-0".to_vec();
let last = b"hist-499".to_vec();
let mut saw_first = false;
let mut saw_last = false;
timeout(Duration::from_secs(5), async {
while !(saw_first && saw_last) {
let f = recv(&mut ws).await;
if f["type"] != "log" {
continue;
}
let msg = f["message"].as_str().unwrap_or("");
if msg.as_bytes() == first.as_slice() {
saw_first = true;
}
if msg.as_bytes() == last.as_slice() {
saw_last = true;
}
}
})
.await
.expect("history replay stalled — handle_subscribe_logs deadlocked");
log_buffer.push_sync("WARN", "after-large-history");
timeout(Duration::from_secs(2), async {
loop {
let f = recv(&mut ws).await;
if f["type"] == "log" && f["message"] == "after-large-history" {
return;
}
}
})
.await
.expect("live entry after large-history replay did not arrive");
}
#[tokio::test]
async fn test_ws_unsubscribe_logs_stops_delivery() {
let (ws_base, _http, log_buffer, _dir) = start_test_server_with_logs().await;
let mut ws = ws_connect(&ws_base).await;
send(&mut ws, json!({"type": "subscribe-logs"})).await;
let _ = recv(&mut ws).await;
send(
&mut ws,
json!({"type": "unsubscribe-logs", "req_id": "ul1"}),
)
.await;
let resp = recv(&mut ws).await;
assert_eq!(resp["type"], "ok");
assert_eq!(resp["req_id"], "ul1");
log_buffer.push_sync("INFO", "after-unsub");
let no_log = timeout(Duration::from_millis(300), async {
loop {
let f = recv(&mut ws).await;
if f["type"] == "log" {
return f;
}
}
})
.await;
assert!(
no_log.is_err(),
"should not receive log frames after unsubscribe"
);
}
#[tokio::test]
async fn test_ws_subscribe_same_queue_twice_is_ok() {
let (ws_base, http_base, _dir) = start_test_server().await;
http_create_queue(&http_base, "dup-sub-q").await;
let mut ws = ws_connect(&ws_base).await;
send(
&mut ws,
json!({"type": "subscribe", "queue": "dup-sub-q", "req_id": "s1"}),
)
.await;
let r1 = recv(&mut ws).await;
assert_eq!(r1["type"], "ok");
send(
&mut ws,
json!({"type": "subscribe", "queue": "dup-sub-q", "req_id": "s2"}),
)
.await;
let r2 = recv(&mut ws).await;
assert_eq!(r2["type"], "ok", "duplicate subscribe should be idempotent");
assert_eq!(r2["req_id"], "s2");
}
#[tokio::test]
async fn test_ws_publish_missing_payload_field() {
let (ws_base, http_base, _dir) = start_test_server().await;
http_create_queue(&http_base, "miss-payload-q").await;
let mut ws = ws_connect(&ws_base).await;
send(
&mut ws,
json!({"type": "publish", "queue": "miss-payload-q", "req_id": "r1"}),
)
.await;
let resp = recv(&mut ws).await;
assert_eq!(resp["type"], "error");
assert_eq!(resp["req_id"], "r1");
}
#[tokio::test]
async fn test_ws_subscribe_missing_queue_field() {
let (ws_base, _http, _dir) = start_test_server().await;
let mut ws = ws_connect(&ws_base).await;
send(&mut ws, json!({"type": "subscribe", "req_id": "r2"})).await;
let resp = recv(&mut ws).await;
assert_eq!(resp["type"], "error");
assert_eq!(resp["req_id"], "r2");
}
#[tokio::test]
async fn test_ws_frame_missing_type_field_returns_error() {
let (ws_base, _http, _dir) = start_test_server().await;
let mut ws = ws_connect(&ws_base).await;
send(&mut ws, json!({"req_id": "r3"})).await;
let resp = recv(&mut ws).await;
assert_eq!(resp["type"], "error");
assert_eq!(resp["req_id"], "r3");
assert_eq!(resp["code"], "invalid_request");
}
#[tokio::test]
async fn test_ws_subscribe_logs_without_buffer_returns_error() {
let (ws_base, _http, _dir) = start_test_server().await; let mut ws = ws_connect(&ws_base).await;
send(
&mut ws,
json!({"type": "subscribe-logs", "req_id": "sl-err"}),
)
.await;
let resp = recv(&mut ws).await;
assert_eq!(resp["type"], "error");
assert_eq!(resp["req_id"], "sl-err");
}
#[tokio::test]
async fn test_ws_subscribe_with_lock_timeout_secs() {
let (ws_base, http_base, _dir) = start_test_server().await;
http_create_queue(&http_base, "long-lock-q").await;
let mut ws = ws_connect(&ws_base).await;
send(
&mut ws,
json!({"type": "subscribe", "queue": "long-lock-q", "lock_timeout_secs": 120}),
)
.await;
let sub_resp = recv(&mut ws).await;
assert_eq!(sub_resp["type"], "ok");
http_publish(&http_base, "long-lock-q", "long-lock-payload").await;
let delivery = timeout(Duration::from_secs(3), async {
loop {
let f = recv(&mut ws).await;
if f["type"] == "deliver" {
return f;
}
}
})
.await
.expect("did not receive delivery frame in time");
let msg_id = delivery["id"].as_str().unwrap().to_string();
let consume_resp: Value = Client::new()
.post(format!("{}/consume/long-lock-q", http_base))
.json(&json!({"consumer_id": "probe", "timeout_seconds": 5}))
.send()
.await
.unwrap()
.json()
.await
.unwrap();
assert!(
consume_resp.get("id").is_none(),
"message should be locked and not consumable by another consumer"
);
let stats: Value = Client::new()
.get(format!("{}/queue-stats/long-lock-q", http_base))
.send()
.await
.unwrap()
.json()
.await
.unwrap();
assert_eq!(
stats["locked"], 1,
"message should be locked after delivery"
);
send(
&mut ws,
json!({"type": "ack", "queue": "long-lock-q", "id": msg_id}),
)
.await;
let ack_resp = recv(&mut ws).await;
assert_eq!(ack_resp["type"], "ok");
}
#[tokio::test]
async fn test_ws_renew_with_lock_timeout_secs() {
let (ws_base, http_base, _dir) = start_test_server().await;
http_create_queue(&http_base, "renew-timeout-q").await;
let mut ws = ws_connect(&ws_base).await;
send(
&mut ws,
json!({"type": "subscribe", "queue": "renew-timeout-q"}),
)
.await;
let _ = recv(&mut ws).await;
http_publish(&http_base, "renew-timeout-q", "renew-timeout-payload").await;
let delivery = timeout(Duration::from_secs(3), async {
loop {
let f = recv(&mut ws).await;
if f["type"] == "deliver" {
return f;
}
}
})
.await
.expect("did not receive delivery frame in time");
let msg_id = delivery["id"].as_str().unwrap().to_string();
send(
&mut ws,
json!({
"type": "renew",
"queue": "renew-timeout-q",
"id": msg_id,
"lock_timeout_secs": 120,
"req_id": "rn-timeout"
}),
)
.await;
let renew_resp = recv(&mut ws).await;
assert_eq!(renew_resp["type"], "ok");
assert_eq!(renew_resp["req_id"], "rn-timeout");
let stats: Value = Client::new()
.get(format!("{}/queue-stats/renew-timeout-q", http_base))
.send()
.await
.unwrap()
.json()
.await
.unwrap();
assert_eq!(
stats["locked"], 1,
"message should still be locked after renew with 120s timeout"
);
send(
&mut ws,
json!({"type": "ack", "queue": "renew-timeout-q", "id": msg_id}),
)
.await;
let ack_resp = recv(&mut ws).await;
assert_eq!(ack_resp["type"], "ok");
}
#[tokio::test]
async fn test_ws_subscribe_default_lock_timeout() {
let (ws_base, http_base, _dir) = start_test_server().await;
http_create_queue(&http_base, "default-lock-q").await;
let mut ws = ws_connect(&ws_base).await;
send(
&mut ws,
json!({"type": "subscribe", "queue": "default-lock-q"}),
)
.await;
let sub_resp = recv(&mut ws).await;
assert_eq!(sub_resp["type"], "ok");
http_publish(&http_base, "default-lock-q", "default-lock-payload").await;
let delivery = timeout(Duration::from_secs(3), async {
loop {
let f = recv(&mut ws).await;
if f["type"] == "deliver" {
return f;
}
}
})
.await
.expect("did not receive delivery frame in time");
let msg_id = delivery["id"].as_str().unwrap();
let stats: Value = Client::new()
.get(format!("{}/queue-stats/default-lock-q", http_base))
.send()
.await
.unwrap()
.json()
.await
.unwrap();
assert_eq!(stats["locked"], 1);
send(
&mut ws,
json!({"type": "ack", "queue": "default-lock-q", "id": msg_id}),
)
.await;
let _ = recv(&mut ws).await;
}
#[tokio::test]
async fn test_ws_outbound_backpressure_does_not_stall_inbound() {
let (ws_base, http_base, _dir) = start_test_server().await;
http_create_queue(&http_base, "bp-sink").await;
http_create_queue(&http_base, "bp-probe").await;
let ws = ws_connect(&ws_base).await;
let (mut ws_tx, mut ws_rx) = ws.split();
let subscribe = json!({"type": "subscribe", "queue": "bp-sink"}).to_string();
ws_tx
.send(TungsteniteMessage::Text(subscribe.into()))
.await
.unwrap();
let drainer = tokio::spawn(async move { while let Some(Ok(_)) = ws_rx.next().await {} });
for i in 0..500 {
http_publish(&http_base, "bp-sink", &format!("flood-{i}")).await;
}
const PROBE_COUNT: u64 = 10;
for i in 0..PROBE_COUNT {
let frame =
json!({"type": "publish", "queue": "bp-probe", "payload": format!("probe-{i}")})
.to_string();
ws_tx
.send(TungsteniteMessage::Text(frame.into()))
.await
.unwrap();
}
let client = Client::new();
let deadline = tokio::time::Instant::now() + Duration::from_secs(10);
let mut total = 0u64;
while tokio::time::Instant::now() < deadline {
let stats: Value = client
.get(format!("{}/queue-stats/bp-probe", http_base))
.send()
.await
.unwrap()
.json()
.await
.unwrap();
total = stats["total"].as_u64().unwrap_or(0);
if total >= PROBE_COUNT {
break;
}
tokio::time::sleep(Duration::from_millis(50)).await;
}
assert_eq!(
total, PROBE_COUNT,
"inbound WS publishes must be processed while the outbound path is busy"
);
drainer.abort();
let _ = ws_tx.close().await;
}