unb-server 2.0.3

mod common;
use common::TestCall;
use std::sync::{Arc, Mutex};

use bytes::Bytes;
use futures_util::{stream, StreamExt};
use serde_json::{json, Value};
use unb::http;
use unb_core::{Envelope, ErrorCode};
use unb_server::dynamic::{self, ErasedCall};
use unb_server::{ContractSchema, HandlerError, Node, OperationContract, ServiceBody};

fn contract(input: Value, output: Option<Value>, event: Option<Value>) -> OperationContract {
    OperationContract {
        input: Some(ContractSchema::Owned(input)),
        output: output.map(ContractSchema::Owned),
        event: event.map(ContractSchema::Owned),
        error: Some(ContractSchema::Owned(json!({ "type": "object" }))),
    }
}

fn body_json<T>(request: &http::Request<T>) -> Value
where
    T: AsRef<[u8]>,
{
    serde_json::from_slice(request.body().as_ref()).unwrap_or(Value::Null)
}

fn unary_call() -> ErasedCall {
    Arc::new(|request| {
        Box::pin(async move {
            Ok(http::Response::builder()
                .body(ServiceBody::Unary(Envelope::encode_payload(&json!({
                    "dynamic": body_json(&request)["value"],
                    "context": request.extensions().get::<String>().cloned()
                }))))
                .unwrap())
        })
    })
}

fn streaming_call() -> ErasedCall {
    Arc::new(|request| {
        Box::pin(async move {
            let event = Envelope::encode_payload(&json!({ "event": body_json(&request)["value"] }));
            Ok(http::Response::builder()
                .body(ServiceBody::Stream(Box::pin(stream::iter(vec![Ok::<
                    Bytes,
                    HandlerError,
                >(
                    event
                )]))))
                .unwrap())
        })
    })
}

#[tokio::test]
async fn runtime_contracts_and_dynamic_dispatch_share_typed_node_behavior() {
    let node = Node::builder("dynamic")
        .layer(dynamic::layer(|mut request, next| async move {
            request.extensions_mut().insert("enriched".to_string());
            next.run(request).await
        }))
        .service(dynamic::unary(
            "probe",
            contract(
                json!({ "runtime": "input" }),
                Some(json!({ "runtime": "output" })),
                None,
            ),
            unary_call(),
        ))
        .service(dynamic::streaming(
            "probe",
            contract(
                json!({ "runtime": "stream-input" }),
                None,
                Some(json!({ "runtime": "event" })),
            ),
            streaming_call(),
        ))
        .insecure_accept_declared_peer_identities()
        .build()
        .unwrap();

    let catalog = node.local_catalog(true);
    let probe = catalog["subjects"]
        .as_array()
        .unwrap()
        .iter()
        .find(|entry| entry["subject"] == "probe")
        .unwrap();
    assert_eq!(
        probe["operations"]["unary"]["input_schema"],
        json!({ "runtime": "input" })
    );
    assert_eq!(
        probe["operations"]["streaming"]["event_schema"],
        json!({ "runtime": "event" })
    );
    assert_eq!(
        node.request("/dynamic/probe", json!({ "value": 7 }))
            .await
            .unwrap(),
        json!({ "dynamic": 7, "context": "enriched" })
    );
    let mut events = node
        .subscribe("/dynamic/probe", json!({ "value": 9 }))
        .await
        .unwrap();
    let event: Value = serde_json::from_slice(&events.next().await.unwrap().unwrap()).unwrap();
    assert_eq!(event, json!({ "event": 9 }));
}

#[tokio::test]
async fn dynamic_layers_preserve_order_rejection_and_consuming_next() {
    let trace = Arc::new(Mutex::new(Vec::new()));
    let first = trace.clone();
    let second = trace.clone();
    let handler = trace.clone();
    let terminal: ErasedCall = Arc::new(move |_| {
        let handler = handler.clone();
        Box::pin(async move {
            handler.lock().unwrap().push("handler");
            Ok(http::Response::builder()
                .body(ServiceBody::Unary(Bytes::new()))
                .unwrap())
        })
    });
    let node = Node::builder("layers")
        .layer(dynamic::layer(move |request, next| {
            let first = first.clone();
            async move {
                first.lock().unwrap().push("first-before");
                let response = next.run(request).await?;
                first.lock().unwrap().push("first-after");
                Ok(response)
            }
        }))
        .layer(dynamic::layer(move |request, next| {
            let second = second.clone();
            async move {
                second.lock().unwrap().push("second");
                if body_json(&request)["reject"] == true {
                    return Err(HandlerError::new(ErrorCode::Unauthorized, "rejected"));
                }
                next.run(request).await
            }
        }))
        .service(dynamic::unary(
            "ordered",
            OperationContract::unknown(),
            terminal,
        ))
        .insecure_accept_declared_peer_identities()
        .build()
        .unwrap();

    node.request("/layers/ordered", json!({})).await.unwrap();
    assert_eq!(
        *trace.lock().unwrap(),
        vec!["first-before", "second", "handler", "first-after"]
    );
    let error = node
        .request("/layers/ordered", json!({ "reject": true }))
        .await
        .unwrap_err();
    assert_eq!(error.code, ErrorCode::Unauthorized);
}