PiCloud/crates/picloud/tests/queue_e2e.rs

//! v1.1.9 queue end-to-end tests.
//!
//! Wires the full all-in-one app via `build_app` (spawns the real
//! dispatcher + queue arm + reclaim task), creates an app + a consumer
//! script + a `queue:receive` trigger, enqueues messages through the
//! admin path, and observes the handler's side effect via a KV marker.
//!
//! ## Gating
//!
//! Skips when `DATABASE_URL` is unset — mirrors `dispatcher_e2e.rs`.
//!
//! ## Handler observation
//!
//! Consumer handlers write `ctx.event` to a KV marker on a collection
//! no trigger watches. Tests poll the marker for the expected attempt
//! count / value.

#![allow(clippy::needless_pass_by_value)]

use std::time::Duration;

use axum_test::TestServer;
use serde_json::{json, Value};
use sqlx::postgres::PgPoolOptions;
use sqlx::PgPool;
use uuid::Uuid;

async fn pool_or_skip() -> Option<PgPool> {
    let Ok(url) = std::env::var("DATABASE_URL") else {
        eprintln!("queue_e2e: DATABASE_URL unset — skipping");
        return None;
    };
    let pool = PgPoolOptions::new()
        .max_connections(5)
        .connect(&url)
        .await
        .expect("connect to DATABASE_URL");
    sqlx::migrate!("../manager-core/migrations")
        .run(&pool)
        .await
        .expect("apply migrations");
    Some(pool)
}

async fn server_for(pool: PgPool, suffix: &str) -> (TestServer, String) {
    use picloud_manager_core::auth::hash_password;
    use picloud_shared::InstanceRole;

    let unique = format!("{suffix}-{}", Uuid::new_v4().simple());
    let auth = picloud::AuthDeps::from_pool(pool.clone());
    let username = format!("e2e-{unique}");
    let hash = hash_password("pw").expect("hash");
    auth.users
        .create(&username, &hash, InstanceRole::Owner, None)
        .await
        .expect("seed admin");

    let app = picloud::build_app(
        pool,
        auth,
        picloud_shared::MasterKey::from_bytes([0x42u8; 32]),
    )
    .await
    .expect("build_app");
    let mut server = TestServer::new(app).expect("TestServer");
    let resp = server
        .post("/api/v1/admin/auth/login")
        .json(&json!({ "username": username, "password": "pw" }))
        .await;
    resp.assert_status_ok();
    let token = resp.json::<Value>()["token"]
        .as_str()
        .expect("login token")
        .to_string();
    server.add_header("authorization", format!("Bearer {token}"));

    let slug = format!("e2e-{unique}");
    let created: Value = server
        .post("/api/v1/admin/apps")
        .json(&json!({ "slug": slug, "name": slug }))
        .await
        .json();
    let app_id = created["id"].as_str().expect("app id").to_string();
    (server, app_id)
}

async fn create_script(server: &TestServer, app_id: &str, name: &str, source: &str) -> String {
    let created: Value = server
        .post("/api/v1/admin/scripts")
        .json(&json!({ "app_id": app_id, "name": name, "source": source }))
        .await
        .json();
    created["id"].as_str().expect("script id").to_string()
}

/// The acker handler records ctx.event into KV and returns successfully.
const ACK_HANDLER: &str = r#"
    let e = ctx.event;
    kv::collection("e2e_markers").set("marker", e);
    #{ ok: true }
"#;

/// The throwing handler records its event and then throws — used to
/// exercise nack → retry → dead-letter.
const THROW_HANDLER: &str = r#"
    let e = ctx.event;
    kv::collection("e2e_markers").set("marker", e);
    throw "boom"
"#;

/// Poll until the queue depth reaches `want` (or give up and return what we saw).
async fn poll_queue_depth(pool: &PgPool, app_id: &str, queue: &str, want: i64) -> i64 {
    let mut depth = -1;
    for _ in 0..200 {
        depth = count_queue_messages(pool, app_id, queue).await;
        if depth == want {
            return depth;
        }
        tokio::time::sleep(Duration::from_millis(50)).await;
    }
    depth
}

async fn poll_marker(pool: &PgPool, app_id: &str) -> Option<Value> {
    for _ in 0..200 {
        let row: Option<(Value,)> = sqlx::query_as(
            "SELECT value FROM kv_entries WHERE app_id = $1 \
             AND collection = 'e2e_markers' AND key = 'marker'",
        )
        .bind(Uuid::parse_str(app_id).expect("uuid"))
        .fetch_optional(pool)
        .await
        .expect("kv read");
        if let Some((v,)) = row {
            return Some(v);
        }
        tokio::time::sleep(Duration::from_millis(100)).await;
    }
    None
}

async fn count_queue_messages(pool: &PgPool, app_id: &str, queue: &str) -> i64 {
    let (n,): (i64,) =
        sqlx::query_as("SELECT COUNT(*) FROM queue_messages WHERE app_id = $1 AND queue_name = $2")
            .bind(Uuid::parse_str(app_id).expect("uuid"))
            .bind(queue)
            .fetch_one(pool)
            .await
            .expect("queue count");
    n
}

async fn count_dead_letters_for_queue(pool: &PgPool, app_id: &str, queue: &str) -> i64 {
    let (n,): (i64,) = sqlx::query_as(
        "SELECT COUNT(*) FROM dead_letters \
         WHERE app_id = $1 AND source = 'queue' \
           AND payload->>'queue_name' = $2",
    )
    .bind(Uuid::parse_str(app_id).expect("uuid"))
    .bind(queue)
    .fetch_one(pool)
    .await
    .expect("dl count");
    n
}

/// Helper: directly INSERT a queue message via SQL (no producer-script
/// path). The producer script path is covered in invoke_e2e + the unit
/// tests on the SDK bridge.
async fn enqueue_directly(pool: &PgPool, app_id: &str, queue: &str, payload: Value) {
    sqlx::query(
        "INSERT INTO queue_messages (app_id, queue_name, payload, max_attempts) \
         VALUES ($1, $2, $3, 3)",
    )
    .bind(Uuid::parse_str(app_id).expect("uuid"))
    .bind(queue)
    .bind(payload)
    .execute(pool)
    .await
    .expect("enqueue");
}

#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn queue_receive_acks_on_success() {
    let Some(pool) = pool_or_skip().await else {
        return;
    };
    let (server, app_id) = server_for(pool.clone(), "ack").await;

    let script_id = create_script(&server, &app_id, "worker", ACK_HANDLER).await;
    let resp = server
        .post(&format!("/api/v1/admin/apps/{app_id}/triggers/queue"))
        .json(&json!({
            "script_id": script_id,
            "queue_name": "jobs",
            "visibility_timeout_secs": 30
        }))
        .await;
    resp.assert_status(axum::http::StatusCode::CREATED);

    enqueue_directly(&pool, &app_id, "jobs", json!({ "x": 42 })).await;

    let marker = poll_marker(&pool, &app_id).await.expect("handler fired");
    assert_eq!(marker["source"], "queue");
    assert_eq!(marker["queue"]["queue_name"], "jobs");
    assert_eq!(marker["queue"]["message"]["x"], 42);

    // Ack deleted the row. The marker is written DURING the handler, but the ack
    // happens after it returns — so poll rather than assert on the instant the
    // marker appears (that window is real, and asserting into it is flaky).
    assert_eq!(poll_queue_depth(&pool, &app_id, "jobs", 0).await, 0);
}

#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn queue_receive_dead_letters_after_max_attempts() {
    let Some(pool) = pool_or_skip().await else {
        return;
    };
    let (server, app_id) = server_for(pool.clone(), "dl").await;

    let script_id = create_script(&server, &app_id, "fail", THROW_HANDLER).await;
    let resp = server
        .post(&format!("/api/v1/admin/apps/{app_id}/triggers/queue"))
        .json(&json!({
            "script_id": script_id,
            "queue_name": "failing",
            "visibility_timeout_secs": 30
        }))
        .await;
    resp.assert_status(axum::http::StatusCode::CREATED);

    enqueue_directly(&pool, &app_id, "failing", json!({ "x": 1 })).await;

    // Allow time for attempts to exhaust (3 retries × exponential
    // backoff ≈ 7s in default config; we cap the poll loop generously).
    for _ in 0..150 {
        if count_dead_letters_for_queue(&pool, &app_id, "failing").await > 0 {
            break;
        }
        tokio::time::sleep(Duration::from_millis(200)).await;
    }
    assert!(
        count_dead_letters_for_queue(&pool, &app_id, "failing").await >= 1,
        "expected a dead-letter row for the queue after max attempts"
    );
    assert_eq!(
        count_queue_messages(&pool, &app_id, "failing").await,
        0,
        "queue row should be deleted after dead-letter"
    );
}

/// Audit fix: dead_letter triggers filtered on `source = "queue"` MUST
/// fire when a queue message exhausts its retries. Before the queue arm
/// called `fan_out_dead_letter`, the DL row was written but no handler
/// delivery was enqueued, so registered `dead_letter` handlers sat idle.
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn queue_dead_letter_fans_out_to_dead_letter_handler() {
    let Some(pool) = pool_or_skip().await else {
        return;
    };
    let (server, app_id) = server_for(pool.clone(), "qdlfan").await;

    // DL handler writes to a different KV key so we don't race the
    // failing handler's own marker write.
    let dl_handler_source = r#"
        kv::collection("e2e_markers").set("dl_marker", ctx.event);
        #{ ok: true }
    "#;
    let dl_handler = create_script(&server, &app_id, "dl-handler", dl_handler_source).await;
    server
        .post(&format!("/api/v1/admin/apps/{app_id}/triggers/dead_letter"))
        .json(&json!({ "script_id": dl_handler, "source_filter": "queue" }))
        .await
        .assert_status(axum::http::StatusCode::CREATED);

    // Failing consumer with max_attempts=1 so we exhaust quickly.
    let failing = create_script(&server, &app_id, "failing", THROW_HANDLER).await;
    server
        .post(&format!("/api/v1/admin/apps/{app_id}/triggers/queue"))
        .json(&json!({
            "script_id": failing,
            "queue_name": "failing-fanout",
            "visibility_timeout_secs": 30,
            "max_attempts": 1
        }))
        .await
        .assert_status(axum::http::StatusCode::CREATED);

    enqueue_directly(&pool, &app_id, "failing-fanout", json!({ "x": 1 })).await;

    // Poll the DL handler's marker.
    for _ in 0..200 {
        let row: Option<(Value,)> = sqlx::query_as(
            "SELECT value FROM kv_entries WHERE app_id = $1 \
             AND collection = 'e2e_markers' AND key = 'dl_marker'",
        )
        .bind(Uuid::parse_str(&app_id).expect("uuid"))
        .fetch_optional(&pool)
        .await
        .expect("kv read");
        if let Some((event,)) = row {
            // Sanity: the dead-letter event names the queue source and
            // nests the original Queue event verbatim.
            assert_eq!(event["source"], "dead_letter");
            assert_eq!(event["dead_letter"]["original"]["source"], "queue");
            assert_eq!(
                event["dead_letter"]["original"]["queue"]["queue_name"],
                "failing-fanout"
            );
            return;
        }
        tokio::time::sleep(Duration::from_millis(100)).await;
    }
    panic!("dead_letter handler never fired for queue-exhausted message");
}

#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn queue_one_consumer_per_queue_rejected() {
    let Some(pool) = pool_or_skip().await else {
        return;
    };
    let (server, app_id) = server_for(pool.clone(), "onecon").await;

    let s1 = create_script(&server, &app_id, "first", ACK_HANDLER).await;
    let s2 = create_script(&server, &app_id, "second", ACK_HANDLER).await;
    let resp = server
        .post(&format!("/api/v1/admin/apps/{app_id}/triggers/queue"))
        .json(&json!({
            "script_id": s1,
            "queue_name": "single_owner",
            "visibility_timeout_secs": 30
        }))
        .await;
    resp.assert_status(axum::http::StatusCode::CREATED);

    let resp2 = server
        .post(&format!("/api/v1/admin/apps/{app_id}/triggers/queue"))
        .json(&json!({
            "script_id": s2,
            "queue_name": "single_owner",
            "visibility_timeout_secs": 30
        }))
        .await;
    // The repo returns Invalid (422) for the duplicate consumer.
    assert!(
        resp2.status_code().is_client_error(),
        "expected 4xx for duplicate consumer, got {}",
        resp2.status_code()
    );
    let body = resp2.text();
    assert!(
        body.contains("already has a consumer")
            || body.contains("queue 'single_owner'")
            || body.contains("consumer trigger"),
        "expected friendly conflict message, got: {body}"
    );
}

#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn queue_visibility_timeout_reclaim() {
    let Some(pool) = pool_or_skip().await else {
        return;
    };
    let (server, app_id) = server_for(pool.clone(), "vt").await;

    let script_id = create_script(&server, &app_id, "slow", ACK_HANDLER).await;
    // Use the API's minimum visibility timeout (30s). The test below
    // forces a stale claim that's 60s old, so it's reclaimable regardless.
    let resp = server
        .post(&format!("/api/v1/admin/apps/{app_id}/triggers/queue"))
        .json(&json!({
            "script_id": script_id,
            "queue_name": "vt_queue",
            "visibility_timeout_secs": 30
        }))
        .await;
    resp.assert_status(axum::http::StatusCode::CREATED);

    // Insert a message and immediately fake a stale claim (claimed_at
    // far enough in the past that the reclaim task picks it up).
    let msg_id = Uuid::new_v4();
    sqlx::query(
        "INSERT INTO queue_messages (id, app_id, queue_name, payload, max_attempts, \
                                     claim_token, claimed_at) \
         VALUES ($1, $2, 'vt_queue', $3, 3, $4, NOW() - INTERVAL '60 seconds')",
    )
    .bind(msg_id)
    .bind(Uuid::parse_str(&app_id).expect("uuid"))
    .bind(json!({ "x": 99 }))
    .bind(Uuid::new_v4())
    .execute(&pool)
    .await
    .expect("insert stale claim");

    // The reclaim task default cadence is 30s; we wait up to 45s for it
    // to fire OR observe the marker (whichever comes first — the marker
    // shows up only after a successful claim by the live dispatcher).
    let marker = poll_marker(&pool, &app_id).await;
    if marker.is_some() {
        return; // dispatcher already re-claimed and delivered
    }
    // Otherwise the reclaim path is the only way the message would
    // have been picked up. Poll the claim state.
    for _ in 0..50 {
        let row: Option<(Option<Uuid>,)> =
            sqlx::query_as("SELECT claim_token FROM queue_messages WHERE id = $1")
                .bind(msg_id)
                .fetch_optional(&pool)
                .await
                .expect("read");
        if let Some((token,)) = row {
            if token.is_none() {
                return; // reclaim succeeded — claim cleared
            }
        } else {
            return; // already delivered + acked
        }
        tokio::time::sleep(Duration::from_secs(1)).await;
    }
    panic!("visibility timeout reclaim did not fire within 50s");
}