maestro/crates/maestro-web/tests/proxy.rs

//! Integration tests for the maestro-web reverse-proxy.
//!
//! Spins up a tiny axum "daemon" on a random port, points maestro-web at it,
//! and exercises both the HTTP and WS proxy paths end-to-end.

use std::net::{Ipv4Addr, SocketAddr};
use std::time::Duration;

use axum::extract::ws::{Message as AxMessage, WebSocket, WebSocketUpgrade};
use axum::response::IntoResponse;
use axum::routing::{any, get};
use axum::{Json, Router};
use futures_util::{SinkExt, StreamExt};
use maestro_web::config::WebConfig;
use tokio::net::TcpListener;
use tokio::sync::oneshot;
use tokio_tungstenite::tungstenite::Message as WsMessage;

async fn spawn_mock_daemon() -> (SocketAddr, oneshot::Sender<()>) {
    async fn list_sessions() -> Json<serde_json::Value> {
        Json(serde_json::json!([{"id": "session-1"}]))
    }
    async fn web_info() -> Json<serde_json::Value> {
        Json(serde_json::json!({"daemon_id": "mock", "hostname": "mock-host"}))
    }
    async fn echo_ws(ws: WebSocketUpgrade) -> impl IntoResponse {
        ws.on_upgrade(echo_handler)
    }
    async fn echo_handler(mut socket: WebSocket) {
        while let Some(Ok(msg)) = socket.next().await {
            match msg {
                AxMessage::Text(t) => {
                    let _ = socket.send(AxMessage::Text(format!("echo:{t}"))).await;
                }
                AxMessage::Binary(b) => {
                    let _ = socket.send(AxMessage::Binary(b)).await;
                }
                AxMessage::Close(_) => break,
                _ => {}
            }
        }
    }

    let app = Router::new()
        .route("/api/sessions", get(list_sessions))
        .route("/api/web/info", get(web_info))
        .route("/api/attach/:id", any(echo_ws));
    let listener = TcpListener::bind(SocketAddr::from((Ipv4Addr::LOCALHOST, 0)))
        .await
        .unwrap();
    let addr = listener.local_addr().unwrap();
    let (tx, rx) = oneshot::channel::<()>();
    tokio::spawn(async move {
        let _ = axum::serve(listener, app)
            .with_graceful_shutdown(async move {
                let _ = rx.await;
            })
            .await;
    });
    (addr, tx)
}

async fn spawn_web(daemon_url: String) -> (SocketAddr, oneshot::Sender<()>) {
    let cfg = WebConfig {
        bind_addr: SocketAddr::from((Ipv4Addr::LOCALHOST, 0)),
        daemon_url,
        dev_proxy_target: None,
    };
    let router = maestro_web::router(&cfg);
    let listener = TcpListener::bind(cfg.bind_addr).await.unwrap();
    let addr = listener.local_addr().unwrap();
    let (tx, rx) = oneshot::channel::<()>();
    tokio::spawn(async move {
        let _ = axum::serve(listener, router)
            .with_graceful_shutdown(async move {
                let _ = rx.await;
            })
            .await;
    });
    (addr, tx)
}

#[tokio::test]
async fn http_get_api_path_forwards_to_daemon() {
    let (daemon_addr, _d) = spawn_mock_daemon().await;
    let (web_addr, _w) = spawn_web(format!("http://{daemon_addr}")).await;

    let body: serde_json::Value = reqwest::get(format!("http://{web_addr}/api/sessions"))
        .await
        .unwrap()
        .json()
        .await
        .unwrap();
    assert_eq!(body[0]["id"], "session-1");
}

#[tokio::test]
async fn http_get_unrelated_path_falls_back_to_spa_index() {
    let (daemon_addr, _d) = spawn_mock_daemon().await;
    let (web_addr, _w) = spawn_web(format!("http://{daemon_addr}")).await;

    // The bundle is empty in test builds — the placeholder index.html serves.
    let resp = reqwest::get(format!("http://{web_addr}/some/route"))
        .await
        .unwrap();
    assert_eq!(resp.status(), reqwest::StatusCode::OK);
    let ct = resp
        .headers()
        .get(reqwest::header::CONTENT_TYPE)
        .and_then(|v| v.to_str().ok())
        .unwrap_or("");
    assert!(ct.starts_with("text/html"), "ct={ct}");
}

#[tokio::test]
async fn ws_api_attach_forwards_to_daemon_and_echoes() {
    let (daemon_addr, _d) = spawn_mock_daemon().await;
    let (web_addr, _w) = spawn_web(format!("http://{daemon_addr}")).await;

    let url = format!("ws://{web_addr}/api/attach/abc");
    let (mut ws, _) = tokio_tungstenite::connect_async(url).await.unwrap();
    ws.send(WsMessage::Text("hello".into())).await.unwrap();

    let echoed = tokio::time::timeout(Duration::from_secs(1), ws.next())
        .await
        .unwrap();
    match echoed {
        Some(Ok(WsMessage::Text(t))) => assert_eq!(t, "echo:hello"),
        other => panic!("unexpected frame: {other:?}"),
    }
}

#[tokio::test]
async fn proxy_returns_502_when_daemon_unreachable() {
    // Bind a port, immediately drop the listener, hand the address to the
    // proxy. Connections to it will be refused.
    let listener = TcpListener::bind(SocketAddr::from((Ipv4Addr::LOCALHOST, 0)))
        .await
        .unwrap();
    let dead_addr = listener.local_addr().unwrap();
    drop(listener);

    let (web_addr, _w) = spawn_web(format!("http://{dead_addr}")).await;

    let resp = reqwest::get(format!("http://{web_addr}/api/sessions"))
        .await
        .unwrap();
    assert_eq!(resp.status(), reqwest::StatusCode::BAD_GATEWAY);
    let body: serde_json::Value = resp.json().await.unwrap();
    assert_eq!(body["error"]["code"], "proxy_failed");
}