aria2-rust-pro/crates/aria2-rust-pro-rpc/src/websocket/tests.rs

use std::{
    collections::BTreeMap,
    sync::{Arc, Mutex},
};

use aria2_rust_pro_core::{
    DownloadEngine, DownloadId, EventListener, RuntimeEvent, RuntimeEventKind,
};

use super::notification::{
    MAX_PENDING_WEBSOCKET_BROADCAST_FRAMES, MAX_PENDING_WEBSOCKET_FRAMES_PER_SESSION,
};
use super::*;
use crate::{InProcessRpcDispatcher, JsonRpcRequest, RpcValue};

#[test]
/// Verifies that notification kinds keep the upstream aria2 WebSocket method names.
fn notification_kind_uses_upstream_websocket_method_names() {
    assert_eq!(
        RpcNotificationKind::DownloadStarted.method_name(),
        "aria2.onDownloadStart"
    );
    assert_eq!(
        RpcNotificationKind::DownloadBtDownloadComplete.method_name(),
        "aria2.onBtDownloadComplete"
    );
}

#[test]
/// Verifies that notification events render the upstream `JSON-RPC` notification shape.
fn websocket_event_renders_upstream_jsonrpc_notification_shape() {
    let event = RpcNotificationEvent {
        kind: RpcNotificationKind::DownloadComplete,
        method: String::new(),
        gid: Some("a1b2c3".to_owned()),
        payload: None,
        meta: crate::model::RpcMeta::default(),
    };

    assert_eq!(
        event.to_websocket_json(),
        "{\"jsonrpc\":\"2.0\",\"method\":\"aria2.onDownloadComplete\",\"params\":[{\"gid\":\"a1b2c3\"}]}"
    );
    assert_eq!(
        event.to_websocket_frame(),
        RpcWebSocketFrame::Text(
            "{\"jsonrpc\":\"2.0\",\"method\":\"aria2.onDownloadComplete\",\"params\":[{\"gid\":\"a1b2c3\"}]}".to_owned()
        )
    );
}

#[test]
/// Verifies that object payloads merge with the gid field in notification output.
fn websocket_event_merges_gid_with_object_payload() {
    let event = RpcNotificationEvent {
        kind: RpcNotificationKind::DownloadError,
        method: String::new(),
        gid: Some("deadbeef".to_owned()),
        payload: Some(RpcValue::Object(BTreeMap::from([(
            "status".to_owned(),
            RpcValue::String("error".to_owned()),
        )]))),
        meta: crate::model::RpcMeta::default(),
    };

    assert_eq!(
        event.to_websocket_json(),
        "{\"jsonrpc\":\"2.0\",\"method\":\"aria2.onDownloadError\",\"params\":[{\"gid\":\"deadbeef\",\"status\":\"error\"}]}"
    );
}

#[test]
/// Verifies that one notification frame is emitted for each matching subscription.
fn registry_emits_one_frame_per_matching_subscription() {
    let mut registry = WebSocketNotificationRegistry::default();
    registry.subscribe(WebSocketSubscription {
        id: "sub-a".to_owned(),
        kind: RpcNotificationKind::DownloadStarted,
        topic: "aria2.onDownloadStart".to_owned(),
    });
    registry.subscribe(WebSocketSubscription {
        id: "sub-b".to_owned(),
        kind: RpcNotificationKind::DownloadStarted,
        topic: "aria2.onDownloadStart".to_owned(),
    });
    registry.subscribe(WebSocketSubscription {
        id: "sub-c".to_owned(),
        kind: RpcNotificationKind::DownloadComplete,
        topic: "aria2.onDownloadComplete".to_owned(),
    });

    let event = RpcNotificationEvent {
        kind: RpcNotificationKind::DownloadStarted,
        method: String::new(),
        gid: Some("feedface".to_owned()),
        payload: None,
        meta: crate::model::RpcMeta::default(),
    };

    let frames = registry.frames_for_event(&event);
    assert_eq!(frames.len(), 2);
    assert!(frames.iter().any(|(id, _)| id == "sub-a"));
    assert!(frames.iter().any(|(id, _)| id == "sub-b"));
    assert!(
        frames
            .iter()
            .all(|(_, frame)| matches!(frame, RpcWebSocketFrame::Text(_)))
    );
}

#[test]
/// Verifies that broadcast queueing fans notification frames out to every connected session.
fn session_registry_broadcasts_notification_frames_to_all_sessions() {
    let mut registry = WebSocketSessionRegistry::default();
    registry.connect("sess-a");
    registry.connect("sess-b");

    let event = RpcNotificationEvent {
        kind: RpcNotificationKind::DownloadComplete,
        method: String::new(),
        gid: Some("abc123".to_owned()),
        payload: None,
        meta: crate::model::RpcMeta::default(),
    };

    registry.queue_event_for_all(&event);

    assert_eq!(registry.pending_count("sess-a"), Some(1));
    assert_eq!(registry.pending_count("sess-b"), Some(1));
    assert_eq!(
        registry.pop_frame("sess-a"),
        Some(RpcWebSocketFrame::Text(
            "{\"jsonrpc\":\"2.0\",\"method\":\"aria2.onDownloadComplete\",\"params\":[{\"gid\":\"abc123\"}]}".to_owned()
        ))
    );
    assert_eq!(
        registry.pop_frame("sess-b"),
        Some(RpcWebSocketFrame::Text(
            "{\"jsonrpc\":\"2.0\",\"method\":\"aria2.onDownloadComplete\",\"params\":[{\"gid\":\"abc123\"}]}".to_owned()
        ))
    );
}

#[test]
/// Verifies that bridged broadcast ingress stays bounded and drops the oldest frame first.
fn session_registry_enforces_bounded_broadcast_ingress_backpressure() {
    let mut registry = WebSocketSessionRegistry::default();
    registry.connect("sess-a");

    for index in 0..(MAX_PENDING_WEBSOCKET_BROADCAST_FRAMES + 8) {
        registry.queue_broadcast_frame(RpcWebSocketFrame::Ping(vec![
            u8::try_from(index % 256).expect("byte should fit"),
        ]));
    }

    assert_eq!(
        registry.pending_broadcast_count(),
        MAX_PENDING_WEBSOCKET_BROADCAST_FRAMES
    );
    assert_eq!(
        registry.pop_frame("sess-a"),
        Some(RpcWebSocketFrame::Ping(vec![8])),
        "oldest bridged broadcast frames should be dropped first once the ingress limit is reached"
    );
}

#[test]
/// Verifies that bridged broadcast ingress fans out only when a session drain occurs.
fn session_registry_drains_broadcast_ingress_into_connected_sessions() {
    let mut registry = WebSocketSessionRegistry::default();
    registry.connect("sess-a");
    registry.connect("sess-b");

    registry.queue_broadcast_frame(RpcWebSocketFrame::Ping(vec![4, 2]));

    assert_eq!(registry.pending_count("sess-a"), Some(0));
    assert_eq!(registry.pending_count("sess-b"), Some(0));
    assert_eq!(registry.pending_broadcast_count(), 1);

    let sess_a_frames = registry.drain_session_frames("sess-a");
    assert_eq!(registry.pending_broadcast_count(), 0);
    assert_eq!(sess_a_frames, vec![RpcWebSocketFrame::Ping(vec![4, 2])]);
    assert_eq!(
        registry.pop_frame("sess-b"),
        Some(RpcWebSocketFrame::Ping(vec![4, 2]))
    );
}

#[test]
/// Verifies targeted queueing, missing-session rejection, and disconnect behavior.
fn session_registry_supports_targeted_queue_and_disconnect() {
    let mut registry = WebSocketSessionRegistry::default();
    registry.connect("sess-a");
    registry.connect("sess-b");

    assert!(registry.queue_frame_for_session("sess-a", RpcWebSocketFrame::Ping(vec![1, 2, 3]),));
    assert!(!registry.queue_frame_for_session("sess-missing", RpcWebSocketFrame::Ping(vec![9]),));

    assert_eq!(registry.pending_count("sess-a"), Some(1));
    assert_eq!(registry.pending_count("sess-b"), Some(0));
    assert_eq!(
        registry.pop_frame("sess-a"),
        Some(RpcWebSocketFrame::Ping(vec![1, 2, 3]))
    );

    registry.disconnect("sess-b");
    assert!(!registry.contains("sess-b"));
    assert_eq!(registry.len(), 1);
}

#[test]
/// Verifies that the per-session queue is bounded and drops the oldest frame under pressure.
fn session_registry_enforces_bounded_pending_frame_backpressure() {
    let mut registry = WebSocketSessionRegistry::default();
    registry.connect("sess-a");

    for index in 0..(MAX_PENDING_WEBSOCKET_FRAMES_PER_SESSION + 8) {
        assert!(registry.queue_frame_for_session(
            "sess-a",
            RpcWebSocketFrame::Ping(vec![u8::try_from(index % 256).expect("byte should fit")]),
        ));
    }

    assert_eq!(
        registry.pending_count("sess-a"),
        Some(MAX_PENDING_WEBSOCKET_FRAMES_PER_SESSION)
    );
    assert_eq!(
        registry.pop_frame("sess-a"),
        Some(RpcWebSocketFrame::Ping(vec![8])),
        "oldest queued frames should be dropped first once the per-session limit is reached"
    );
}

#[test]
/// Verifies that the runtime-event bridge queues real completion events into session frames.
fn runtime_event_bridge_queues_real_download_events_into_sessions() {
    let sessions = Arc::new(Mutex::new(WebSocketSessionRegistry::default()));
    sessions
        .lock()
        .expect("sessions lock should succeed")
        .connect("sess-a");
    let mut bridge = RuntimeEventWebSocketBridge::new(Arc::clone(&sessions));

    bridge.on_event(
        &RuntimeEvent::new(RuntimeEventKind::DownloadCompleted).with_gid(DownloadId::new(0x2a)),
    );

    assert_eq!(
        sessions
            .lock()
            .expect("sessions lock should succeed")
            .pending_broadcast_count(),
        1
    );
    let frame = sessions
        .lock()
        .expect("sessions lock should succeed")
        .pop_frame("sess-a")
        .expect("download completion should queue a notification frame");
    match frame {
        RpcWebSocketFrame::Text(text) => {
            assert!(text.contains("aria2.onDownloadComplete"));
            assert!(text.contains(r#""gid":"000000000000002a""#));
        }
        other => panic!("unexpected bridged notification frame: {other:?}"),
    }
}

#[test]
/// Verifies that `addUri` registration alone does not emit a start notification.
fn runtime_event_bridge_does_not_treat_add_uri_as_download_start() {
    let sessions = Arc::new(Mutex::new(WebSocketSessionRegistry::default()));
    sessions
        .lock()
        .expect("sessions lock should succeed")
        .connect("sess-a");
    let bridge = RuntimeEventWebSocketBridge::new(Arc::clone(&sessions));
    let mut dispatcher = InProcessRpcDispatcher::new();
    dispatcher.register_runtime_listener(bridge);

    let response = dispatcher.dispatch_json(JsonRpcRequest {
        jsonrpc: Some("2.0".to_owned()),
        id: Some(crate::jsonrpc::JsonRpcId::Number(1)),
        method: "aria2.addUri".to_owned(),
        params: vec![RpcValue::String("https://example.org/file.iso".to_owned())],
        meta: crate::model::RpcMeta::default(),
    });
    assert!(response.error.is_none());

    assert!(
        sessions
            .lock()
            .expect("sessions lock should succeed")
            .pop_frame("sess-a")
            .is_none(),
        "addUri should not emit aria2.onDownloadStart before the download actually starts"
    );
}

#[test]
/// Verifies that start notifications only appear after the scheduler activates the download.
fn runtime_event_bridge_emits_start_only_when_scheduler_activates_download() {
    let sessions = Arc::new(Mutex::new(WebSocketSessionRegistry::default()));
    sessions
        .lock()
        .expect("sessions lock should succeed")
        .connect("sess-a");

    let mut engine = DownloadEngine::new();
    engine.register_listener(RuntimeEventWebSocketBridge::new(Arc::clone(&sessions)));
    let gid = engine.add_uri("https://example.org/file.iso").gid();

    assert!(
        sessions
            .lock()
            .expect("sessions lock should succeed")
            .pop_frame("sess-a")
            .is_none(),
        "registration alone should not emit a start notification"
    );

    let _ = engine.schedule_once();

    let frame = sessions
        .lock()
        .expect("sessions lock should succeed")
        .pop_frame("sess-a")
        .expect("scheduler activation should emit start notification");
    match frame {
        RpcWebSocketFrame::Text(text) => {
            assert!(text.contains("aria2.onDownloadStart"));
            assert!(text.contains(&format!(r#""gid":"{gid}""#)));
        }
        other => panic!("unexpected frame: {other:?}"),
    }
}

#[test]
/// Verifies that resumed downloads emit start only after they become active again.
fn runtime_event_bridge_does_not_emit_start_until_resumed_download_is_active_again() {
    let sessions = Arc::new(Mutex::new(WebSocketSessionRegistry::default()));
    sessions
        .lock()
        .expect("sessions lock should succeed")
        .connect("sess-a");

    let mut engine = DownloadEngine::new();
    engine.register_listener(RuntimeEventWebSocketBridge::new(Arc::clone(&sessions)));
    let gid = engine.add_uri("https://example.org/file.iso").gid();
    let _ = engine.schedule_once();

    let _ = sessions
        .lock()
        .expect("sessions lock should succeed")
        .pop_frame("sess-a");

    engine.pause(gid).expect("pause should succeed");
    let pause = sessions
        .lock()
        .expect("sessions lock should succeed")
        .pop_frame("sess-a")
        .expect("pause should emit pause notification");
    match pause {
        RpcWebSocketFrame::Text(text) => {
            assert!(text.contains("aria2.onDownloadPause"));
        }
        other => panic!("unexpected pause frame: {other:?}"),
    }

    engine
        .resume(gid)
        .expect("resume should move download back to waiting");
    assert!(
        sessions
            .lock()
            .expect("sessions lock should succeed")
            .pop_frame("sess-a")
            .is_none(),
        "resume should not emit start notification before reactivation"
    );

    let _ = engine.schedule_once();

    let frame = sessions
        .lock()
        .expect("sessions lock should succeed")
        .pop_frame("sess-a")
        .expect("reactivation should emit start notification");
    match frame {
        RpcWebSocketFrame::Text(text) => {
            assert!(text.contains("aria2.onDownloadStart"));
            assert!(text.contains(&format!(r#""gid":"{gid}""#)));
        }
        other => panic!("unexpected reactivation frame: {other:?}"),
    }
}