nordisk/edda

crates/nornir-flight/tests/roundtrip.rs

//! Acceptance for the transport seam (design §f1):
//! * round-trip — `do_put` N `RecordBatch`es → `do_get` / `do_exchange` returns
//!   them byte-identical;
//! * auth via the shared `nornir-service` interceptor — missing token rejected,
//!   correct accepted, empty = open.
//!
//! Each assertion is RED-when-broken: a transport that drops/mangles batches, or
//! an auth seam that lets an unauthenticated call through, fails here.

use std::net::{SocketAddr, TcpListener};
use std::sync::Arc;
use std::time::Duration;

use arrow_array::{Int64Array, RecordBatch, StringArray};
use arrow_schema::{DataType, Field, Schema};
use nornir_flight::{FlightConn, FlightRoutes, FlightServer, MemFlightStore};

/// Grab an OS-assigned free port (bind, read, drop — the server rebinds it).
fn free_addr() -> SocketAddr {
    TcpListener::bind("127.0.0.1:0")
        .unwrap()
        .local_addr()
        .unwrap()
}

/// `n` distinct two-column batches (Int64 + nullable Utf8).
fn sample_batches(n: usize) -> Vec<RecordBatch> {
    let schema = Arc::new(Schema::new(vec![
        Field::new("id", DataType::Int64, false),
        Field::new("name", DataType::Utf8, true),
    ]));
    (0..n)
        .map(|i| {
            let base = i as i64 * 10;
            let ids = Int64Array::from(vec![base, base + 1, base + 2]);
            let names = StringArray::from(vec![Some(format!("r{i}-a")), None, Some(format!("r{i}-c"))]);
            RecordBatch::try_new(schema.clone(), vec![Arc::new(ids), Arc::new(names)]).unwrap()
        })
        .collect()
}

/// Start a server (fresh `MemFlightStore`) on a free port with `token` auth
/// (empty ⇒ open), and return its `http://…` endpoint once it accepts connects.
async fn start_server(token: &str) -> String {
    let addr = free_addr();
    let routes = FlightRoutes::shared(Arc::new(MemFlightStore::new()));
    let tok = Arc::new(token.to_string());
    tokio::spawn(async move {
        let _ = FlightServer::serve_with_auth(routes, addr, tok).await;
    });
    for _ in 0..100 {
        if tokio::net::TcpStream::connect(addr).await.is_ok() {
            break;
        }
        tokio::time::sleep(Duration::from_millis(20)).await;
    }
    format!("http://{addr}")
}

#[tokio::test]
async fn put_then_get_roundtrip_byte_identical() {
    let ep = start_server("").await;
    let mut c = FlightConn::connect(ep).await.unwrap();
    let batches = sample_batches(3);
    c.put("build/42", batches.clone()).await.unwrap();
    let got = c.get("build/42").await.unwrap();
    assert_eq!(
        got, batches,
        "do_get must return the do_put batches byte-identical"
    );
}

#[tokio::test]
async fn exchange_roundtrip_byte_identical() {
    let ep = start_server("").await;
    let mut c = FlightConn::connect(ep).await.unwrap();
    let batches = sample_batches(4);
    let got = c.exchange("x/1", batches.clone()).await.unwrap();
    assert_eq!(got, batches, "do_exchange must echo the batches byte-identical");
}

#[tokio::test]
async fn empty_token_is_open() {
    // Server with empty token = open: a client with NO bearer round-trips fine.
    let ep = start_server("").await;
    let mut c = FlightConn::connect(ep).await.unwrap();
    let batches = sample_batches(2);
    c.put("k", batches.clone()).await.unwrap();
    assert_eq!(c.get("k").await.unwrap(), batches);
}

#[tokio::test]
async fn auth_correct_token_accepted() {
    let token = "a-long-enough-secret-token-1234";
    let ep = start_server(token).await;
    let mut c = FlightConn::connect(ep).await.unwrap().with_bearer(token).unwrap();
    let batches = sample_batches(2);
    c.put("k", batches.clone()).await.unwrap();
    assert_eq!(
        c.get("k").await.unwrap(),
        batches,
        "correct bearer must be accepted and round-trip"
    );
}

#[tokio::test]
async fn auth_missing_token_rejected() {
    let token = "a-long-enough-secret-token-1234";
    let ep = start_server(token).await;
    let mut c = FlightConn::connect(ep).await.unwrap(); // no bearer
    let err = c
        .put("k", sample_batches(1))
        .await
        .expect_err("unauthenticated put MUST be rejected");
    let msg = err.to_string().to_lowercase();
    assert!(
        msg.contains("bearer") || msg.contains("unauthenticated"),
        "expected an auth rejection, got: {msg}"
    );
}

#[tokio::test]
async fn auth_wrong_token_rejected() {
    let token = "a-long-enough-secret-token-1234";
    let ep = start_server(token).await;
    let mut c = FlightConn::connect(ep)
        .await
        .unwrap()
        .with_bearer("the-wrong-token")
        .unwrap();
    c.put("k", sample_batches(1))
        .await
        .expect_err("wrong bearer MUST be rejected");
}

#[tokio::test]
async fn unknown_ticket_is_not_found() {
    let ep = start_server("").await;
    let mut c = FlightConn::connect(ep).await.unwrap();
    c.get("never-put")
        .await
        .expect_err("a get on an unknown key MUST error (NOT_FOUND)");
}