unb-server 2.0.3

#![allow(dead_code)]

use std::sync::Arc;
use std::time::Duration;

use axum::extract::ws::WebSocketUpgrade;
use axum::routing::get;
use axum::Router;
use serde_json::Value;
use unb_core::Envelope;
use unb_runtime::{ClientStream, SessionOutcome, Wire};
use unb_server::{HandlerError, Node, SendExt};

pub trait TestCall {
    fn request(
        &self,
        target_path: &str,
        payload: Value,
    ) -> impl std::future::Future<Output = Result<Value, HandlerError>> + Send;
}

impl TestCall for Arc<Node> {
    async fn request(&self, target_path: &str, payload: Value) -> Result<Value, HandlerError> {
        let response = http::Request::builder()
            .uri(target_path)
            .body(payload)
            .send(self)
            .await?;
        Ok(if response.body().is_empty() {
            Value::Null
        } else {
            serde_json::from_slice(response.body()).unwrap_or(Value::Null)
        })
    }
}

pub async fn start(node: Arc<Node>) -> String {
    let app = Router::new().route(
        "/ws",
        get(move |upgrade: WebSocketUpgrade| {
            let node = node.clone();
            async move { node.serve_ws_upgrade(upgrade) }
        }),
    );
    let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
    let addr = listener.local_addr().unwrap();
    tokio::spawn(async move { axum::serve(listener, app).await.unwrap() });
    format!("ws://{addr}/ws")
}

pub async fn ready_client(wire: &Wire) {
    let outcome = wire.session_outcome().await.unwrap();
    assert!(
        matches!(outcome, SessionOutcome::Established),
        "{outcome:?}"
    );
}

pub async fn connect_nodes(dialer: &Arc<Node>, peer: &str, acceptor: &Arc<Node>) {
    dialer
        .link(acceptor)
        .await
        .unwrap_or_else(|error| panic!("link to {peer:?} failed: {error}"));
}

pub async fn ghost_establish(wire: &Wire, node_id: &str) {
    ghost_establish_with(wire, node_id, &[]).await
}

pub async fn ghost_establish_with(wire: &Wire, node_id: &str, _subjects: &[&str]) {
    use unb_core::{Kind, NodeIdentity, RouteAck, RouteAckStatus, RouteSnapshot};
    let identity = NodeIdentity {
        node_id: node_id.into(),
        instance_id: format!("{node_id}-1"),
        epoch: 1,
        proof: serde_json::Value::Null,
    };
    let mut sent_identify = false;
    let mut sent_snapshot = false;
    let mut got_snapshot = false;
    let mut got_ack = false;
    let mut observer = wire.observe();
    while let Ok(envelope) = observer.recv().await {
        match envelope.kind {
            Kind::Identify => {
                if !sent_identify {
                    let payload =
                        Envelope::encode_payload(&serde_json::to_value(&identity).unwrap());
                    wire.control(Kind::Identify, payload).await.unwrap();
                    sent_identify = true;
                }
                wire.control(
                    Kind::IdentityAccepted,
                    Envelope::encode_payload(&serde_json::Value::Null),
                )
                .await
                .unwrap();
            }
            Kind::IdentityAccepted => {
                if !sent_snapshot {
                    let routes = vec![unb_core::RouteAdvertisement {
                        destination: node_id.to_string(),
                        owner: node_id.to_string(),
                        owner_instance: identity.instance_id.clone(),
                        owner_epoch: identity.epoch,
                        owner_revision: 0,
                        distance: 0,
                        path: vec![node_id.to_string()],
                    }];
                    let snapshot = RouteSnapshot::canonical(1, routes);
                    let payload =
                        Envelope::encode_payload(&serde_json::to_value(&snapshot).unwrap());
                    wire.control(Kind::RouteSnapshot, payload).await.unwrap();
                    sent_snapshot = true;
                }
            }
            Kind::RouteSnapshot => {
                let generation = envelope.payload_json()["generation"].as_u64().unwrap_or(1);
                let ack = RouteAck {
                    generation,
                    status: RouteAckStatus::Applied,
                };
                let payload = Envelope::encode_payload(&serde_json::to_value(&ack).unwrap());
                wire.control(Kind::RouteAck, payload).await.unwrap();
                got_snapshot = true;
            }
            Kind::RouteAck => got_ack = true,
            _ => {}
        }
        if sent_identify && sent_snapshot && got_snapshot && got_ack {
            return;
        }
    }
}

pub async fn wait_until(label: &str, cond: impl Fn() -> bool) {
    tokio::time::timeout(Duration::from_secs(5), async {
        while !cond() {
            tokio::time::sleep(Duration::from_millis(10)).await;
        }
    })
    .await
    .unwrap_or_else(|_| panic!("condition never held: {label}"));
}

pub async fn deliver(wire: &Wire) -> Envelope {
    deliver_observed(&mut wire.observe()).await
}

pub async fn deliver_observed(
    observer: &mut tokio::sync::broadcast::Receiver<Envelope>,
) -> Envelope {
    loop {
        let envelope = observer
            .recv()
            .await
            .expect("wire closed while waiting for a frame");
        if matches!(
            envelope.kind,
            unb_core::Kind::Identify
                | unb_core::Kind::IdentityAccepted
                | unb_core::Kind::RouteSnapshot
                | unb_core::Kind::RouteDelta
                | unb_core::Kind::RouteAck
        ) {
            continue;
        }
        return envelope;
    }
}

pub async fn stream(
    wire: &Wire,
    subject: &str,
    kind: unb_core::Kind,
    payload: Value,
) -> ClientStream {
    ready_client(wire).await;
    wire.client_session()
        .start(
            subject,
            kind,
            Envelope::encode_payload(&payload),
            None,
            Default::default(),
        )
        .await
        .unwrap()
}