continuum-proxy/tests/uplink_outbox.rs

//! Job outcomes survive the cloud being away: an event in the gateway's store
//! is sent on connect, sent again after a reconnect if nobody acknowledged it,
//! and gone once acknowledged. A start that fails after being accepted is
//! queued the same way instead of being sent once to whoever is listening.
//!
//! Real `uplink::run` against a real websocket server; no printer is touched
//! (the fleet is empty, so the start fails on its lookup before any download).

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

use futures_util::{SinkExt, StreamExt};
use serde_json::{json, Value};
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::TcpListener;
use tokio::sync::Mutex;
use tokio_tungstenite::WebSocketStream;
use tokio_tungstenite::tungstenite::Message;
use tokio_tungstenite::tungstenite::handshake::server::{Request, Response};

use continuum_proxy::config::Config;
use continuum_proxy::store::Store;
use continuum_proxy::uplink;

/// Long enough for a loopback round trip, short enough that a regression
/// fails the test instead of hanging the suite.
const WAIT: Duration = Duration::from_secs(10);

fn config(port: u16) -> Config {
    Config {
        gateway_id: "test-gw".into(),
        farm_id: "test-farm".into(),
        gateway_token: "test-token".into(),
        uplink_url: format!("ws://127.0.0.1:{port}/ws/edge/v1"),
        // Nothing here starts go2rtc or reads the fleet file - `uplink::run`
        // is handed its fleet rather than loading one - but Config has no
        // Default, so these have to say something.
        go2rtc_bin: "/nonexistent-go2rtc".into(),
        go2rtc_config: "/nonexistent/go2rtc.yaml".into(),
        go2rtc_api_listen: "127.0.0.1:1984".into(),
        stream_base_url: None,
        printers_path: "/nonexistent/printers.toml".into(),
        store_path: "/nonexistent/continuum.db".into(),
        // Off, which is the default and also the only safe thing in a test:
        // an enabled one would sweep whatever network the test host is on.
        discovery: Default::default(),
    }
}

/// The next frame as JSON. Heartbeats and telemetry share this socket, so
/// callers filter; `interval` fires on its first tick, which means a
/// heartbeat arrives immediately after the hello.
async fn next_json<S>(ws: &mut WebSocketStream<S>) -> Value
where
    S: tokio::io::AsyncRead + tokio::io::AsyncWrite + Unpin,
{
    loop {
        let frame = tokio::time::timeout(WAIT, ws.next())
            .await
            .expect("timed out waiting for a frame from the gateway")
            .expect("the gateway closed the connection")
            .expect("websocket error");
        if let Message::Text(text) = frame {
            return serde_json::from_str(&text).expect("the gateway sent invalid JSON");
        }
    }
}

/// Plays the cloud's `/auth/edge-token`: reads one HTTP request, checks that
/// the gateway sent its farm and long-lived token, and answers with a JWT.
/// Returns the request text so the caller can assert on it.
async fn answer_token_exchange(listener: &TcpListener) -> String {
    let (mut stream, _) = listener.accept().await.unwrap();
    let mut request = Vec::new();
    let mut buf = [0u8; 4096];
    loop {
        let n = tokio::time::timeout(WAIT, stream.read(&mut buf)).await.unwrap().unwrap();
        assert!(n > 0, "the gateway hung up before finishing its token request");
        request.extend_from_slice(&buf[..n]);
        let text = String::from_utf8_lossy(&request).into_owned();
        if let Some((head, body)) = text.split_once("\r\n\r\n") {
            let length = head
                .lines()
                .find_map(|l| l.to_ascii_lowercase().strip_prefix("content-length:").map(|v| v.trim().parse::<usize>().unwrap()))
                .unwrap_or(0);
            if body.len() >= length {
                break;
            }
        }
    }
    let body = r#"{"token":"minted-edge-jwt","expiresInSeconds":3600}"#;
    stream
        .write_all(format!("HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{body}", body.len()).as_bytes())
        .await
        .unwrap();
    String::from_utf8_lossy(&request).into_owned()
}


/// Answers the token exchange and accepts the websocket, returning it after
/// the gateway's hello.
async fn accept(listener: &TcpListener) -> WebSocketStream<tokio::net::TcpStream> {
    answer_token_exchange(listener).await;
    let (stream, _) = listener.accept().await.unwrap();
    let mut ws = tokio_tungstenite::accept_hdr_async(stream, |_: &Request, response: Response| Ok(response)).await.unwrap();
    let hello = next_json(&mut ws).await;
    assert_eq!(hello["type"], "hello");
    ws
}

/// Reads frames until `want` job events have arrived, returning them by job id.
async fn job_events(ws: &mut WebSocketStream<tokio::net::TcpStream>, want: usize) -> Vec<Value> {
    let mut events = Vec::new();
    while events.len() < want {
        let msg = next_json(ws).await;
        if msg["type"] == "job_event" {
            events.push(msg);
        }
    }
    events.sort_by_key(|e| e["event_id"].as_i64());
    events
}

#[tokio::test]
async fn unacknowledged_job_events_are_resent_after_a_reconnect_and_stop_once_acknowledged() {
    let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
    let port = listener.local_addr().unwrap().port();

    let dir = std::env::temp_dir().join(format!("continuum-outbox-test-{}", std::process::id()));
    let store = Arc::new(Store::open(&dir.join("s.db")).unwrap());
    // An outcome reached while the cloud was away.
    store.outbox_push("11111111-1111-4111-8111-111111111111", "finished", None, 1_000).unwrap();

    let fleet: uplink::SharedFleet = Arc::new(Mutex::new(Vec::new()));
    let (outbound, inbox) = uplink::channel();
    tokio::spawn(uplink::run(config(port), fleet, Some(store.clone()), outbound, inbox));

    let mut ws = accept(&listener).await;
    let first = job_events(&mut ws, 1).await;
    assert_eq!(first[0]["job_id"], "11111111-1111-4111-8111-111111111111");
    assert_eq!(first[0]["event"], "finished");
    assert_eq!(first[0]["at_ms"], 1_000);

    // A start that fails after being accepted goes through the same outbox.
    ws.send(Message::Text(
        json!({
            "type": "command", "command_id": "cmd-1", "printer_id": "not-in-this-fleet",
            "command": { "action": "start_print", "job_id": "22222222-2222-4222-8222-222222222222",
                         "file_name": "cube.gcode", "url": "http://127.0.0.1:1/never-fetched" }
        })
        .to_string(),
    ))
    .await
    .unwrap();
    // Only the new one: the first was already sent on this connection.
    let mut both = first.clone();
    both.extend(job_events(&mut ws, 1).await);
    assert_eq!(both[1]["job_id"], "22222222-2222-4222-8222-222222222222");
    assert_eq!(both[1]["event"], "failed");
    assert!(both[1]["detail"].as_str().unwrap().contains("not-in-this-fleet"), "{}", both[1]);

    // The cloud never answers; the connection drops.
    drop(ws);

    let mut ws = accept(&listener).await;
    let again = job_events(&mut ws, 2).await;
    assert_eq!(
        again.iter().map(|e| e["event_id"].clone()).collect::<Vec<_>>(),
        both.iter().map(|e| e["event_id"].clone()).collect::<Vec<_>>(),
        "after a reconnect everything unacknowledged is sent again"
    );

    ws.send(Message::Text(json!({ "type": "job_event_ack", "event_id": again[0]["event_id"] }).to_string()))
        .await
        .unwrap();
    let deadline = tokio::time::Instant::now() + WAIT;
    while store.outbox_pending().unwrap().len() != 1 {
        assert!(tokio::time::Instant::now() < deadline, "the acknowledged event was not cleared");
        tokio::time::sleep(Duration::from_millis(50)).await;
    }
    assert_eq!(store.outbox_pending().unwrap()[0].job_id, "22222222-2222-4222-8222-222222222222");
    let _ = std::fs::remove_dir_all(dir);
}