unb-server 2.0.3

mod common;

use std::time::Duration;

use schemars::JsonSchema;
use serde::{Deserialize, Serialize};
use serde_json::json;
use unb::{handler, HandlerError, Reply, Request};
use unb_client::pair;
use unb_core::{Envelope, Kind, PROTOCOL_VERSION};
use unb_runtime::{Pipe, Wire};
use unb_server::Node;

#[derive(Deserialize, Serialize, JsonSchema)]
#[serde(transparent)]
struct EchoPayload(serde_json::Value);

#[handler]
async fn echo(request: Request<EchoPayload>) -> Result<Reply<EchoPayload>, HandlerError> {
    Ok(Reply::new(request.into_payload()))
}

fn hello() -> Envelope {
    Envelope {
        v: PROTOCOL_VERSION,
        id: "f1".into(),
        target: String::new(),
        subject: String::new(),
        kind: Kind::Hello,
        corr: None,
        seq: None,
        hops: None,
        body_token: None,
        payload: Envelope::encode_payload(&json!({ "versions": [PROTOCOL_VERSION] })),
        path: Vec::new(),
        headers: Default::default(),
    }
}

fn classify_request() -> Envelope {
    Envelope {
        v: PROTOCOL_VERSION,
        id: "f2".into(),
        target: "keeper".into(),
        subject: "echo".into(),
        kind: Kind::Request,
        corr: Some("c1".into()),
        seq: None,
        hops: None,
        body_token: None,
        payload: Envelope::encode_payload(&json!({})),
        path: Vec::new(),
        headers: Default::default(),
    }
}

#[tokio::test(start_paused = true)]
async fn a_silent_peer_is_reaped_by_keepalive() {
    let node = Node::builder("keeper")
        .service(echo)
        .insecure_accept_declared_peer_identities()
        .build()
        .unwrap();
    let (manual_side, node_side) = pair();
    let wire = node.serve_transport(node_side).await;
    let Pipe::Local {
        rx: mut manual_rx,
        tx: manual_tx,
        ..
    } = manual_side
    else {
        panic!("pair() must hand back a local transport");
    };

    manual_tx.send(hello()).await.unwrap();
    let welcome = tokio::time::timeout(Duration::from_secs(5), manual_rx.recv())
        .await
        .expect("the acceptor answers a valid hello")
        .expect("welcome frame");
    assert_eq!(welcome.kind, Kind::Welcome);
    manual_tx.send(classify_request()).await.unwrap();
    let response = tokio::time::timeout(Duration::from_secs(5), manual_rx.recv())
        .await
        .expect("the classified client is answered")
        .expect("response frame");
    assert_eq!(response.kind, Kind::Response);

    tokio::time::advance(Duration::from_secs(46)).await;

    let mut saw_ping = false;
    let closed = tokio::time::timeout(Duration::from_secs(120), async {
        while let Some(frame) = manual_rx.recv().await {
            if frame.kind == Kind::Ping {
                saw_ping = true;
            }
        }
    })
    .await;
    assert!(
        closed.is_ok(),
        "a peer that never answers pings must be reaped within the keepalive window"
    );
    assert!(saw_ping, "the session pings an idle peer before reaping it");

    let dead = tokio::time::timeout(Duration::from_secs(10), async {
        loop {
            if wire.respond("s1", json!({})).await.is_err() {
                return;
            }
            tokio::task::yield_now().await;
        }
    })
    .await;
    assert!(dead.is_ok(), "the reaped session's wire must be gone");
}

#[tokio::test(start_paused = true)]
async fn an_idle_healthy_peer_stays_established() {
    let node = Node::builder("keeper")
        .service(echo)
        .insecure_accept_declared_peer_identities()
        .build()
        .unwrap();
    let (client_side, node_side) = pair();
    node.serve_transport(node_side).await;
    let client = Wire::open(client_side);
    common::ready_client(&client).await;
    let mut call = client
        .client_session()
        .start(
            "/keeper/echo",
            Kind::Request,
            Envelope::encode_payload(&json!({ "warm": true })),
            None,
            Default::default(),
        )
        .await
        .expect("the client classifies with one request");
    tokio::time::timeout(Duration::from_secs(5), call.next())
        .await
        .expect("the classifying call round-trips")
        .unwrap()
        .unwrap();

    tokio::time::advance(Duration::from_secs(120)).await;

    let mut call = client
        .client_session()
        .start(
            "/keeper/echo",
            Kind::Request,
            Envelope::encode_payload(&json!({ "alive": true })),
            None,
            Default::default(),
        )
        .await
        .expect("an idle-but-healthy session must survive keepalive");
    let response = tokio::time::timeout(Duration::from_secs(10), call.next())
        .await
        .expect("the call must round-trip after long idling")
        .unwrap()
        .unwrap();
    assert_eq!(response.kind, Kind::Response);
    assert_eq!(response.payload_json()["alive"], true);
}