maarcadetweet/crates/appview/tests/pds_firehose_integration.rs

//! End-to-end tests for the local PDS firehose consumer.
//!
//! These are the only tests that put *real* PDS bytes through
//! [`appview::pds_firehose`]: everything else in the module's own
//! `#[cfg(test)]` section builds frames from a hand-written DAG-CBOR
//! encoder, which proves the decoder matches our reading of the
//! contract but not that the PDS writes what we think it writes.
//!
//! Fail-open, like every other suite in this directory. Each test
//! prints a notice and returns successfully when a precondition is
//! missing:
//!
//!  - the PDS isn't running on `:2583`;
//!  - `DATABASE_URL_APPVIEW` is unset or the database is unreachable;
//!  - **`com.atproto.sync.subscribeRepos` does not exist yet.** The
//!    endpoint is being built in `crates/pds-server` in parallel with
//!    this consumer. Until it lands, the WebSocket upgrade fails and
//!    these tests skip with a message saying so — they are not proof of
//!    anything while that line appears in the output.
//!
//! What they cover once the endpoint is live:
//!
//!  - `frame_from_the_local_pds_indexes_a_post` — subscribe, create a
//!    record over `com.atproto.repo.createRecord`, and drive the frame
//!    the PDS emits through the real decode → CAR → indexer path,
//!    asserting the row lands in `posts`.
//!  - `replaying_the_same_frame_changes_nothing` — the same frame
//!    applied twice leaves exactly one row, which is what makes the
//!    overlap with the `/internal/ingest-commit` push safe.
//!  - `car_reader_parses_a_real_repo_export` — the AppView's CAR reader
//!    against a CAR the PDS's *writer* produced (`getRepo`).
//!  - `healthz_reports_the_firehose_state` — the running AppView's
//!    probe carries the new fields.

use futures::StreamExt;
use serde_json::{json, Value};
use sqlx::PgPool;
use std::time::Duration;
use tokio_tungstenite::tungstenite::Message;

use appview::pds_firehose::{self, Frame};

const PDS_URL: &str = "http://127.0.0.1:2583";
const PDS_WS: &str = "ws://127.0.0.1:2583";

fn appview_url() -> String {
    std::env::var("APPVIEW_TEST_URL").unwrap_or_else(|_| "http://127.0.0.1:2584".to_string())
}

fn client() -> reqwest::Client {
    reqwest::Client::builder()
        .timeout(Duration::from_secs(10))
        .build()
        .unwrap()
}

async fn service_up(c: &reqwest::Client, base: &str) -> bool {
    for _ in 0..12 {
        if let Ok(r) = c.get(format!("{base}/healthz")).send().await {
            if r.status().is_success() {
                return true;
            }
        }
        tokio::time::sleep(Duration::from_millis(250)).await;
    }
    false
}

async fn appview_db() -> Option<PgPool> {
    let url = std::env::var("DATABASE_URL_APPVIEW").ok()?;
    match tokio::time::timeout(Duration::from_secs(2), PgPool::connect(&url)).await {
        Ok(Ok(pool)) => Some(pool),
        _ => None,
    }
}

struct Account {
    did: String,
    access_jwt: String,
}

async fn create_account(c: &reqwest::Client) -> Option<Account> {
    let handle = format!("fh_{}.maarcadetweet.local", uuid::Uuid::new_v4().simple());
    let r: Value = c
        .post(format!("{PDS_URL}/xrpc/com.atproto.server.createAccount"))
        .json(&json!({ "handle": handle, "password": "hunter2hunter2" }))
        .send()
        .await
        .ok()?
        .json()
        .await
        .ok()?;
    Some(Account {
        did: r["did"].as_str()?.to_string(),
        access_jwt: r["access_jwt"].as_str()?.to_string(),
    })
}

async fn create_post(c: &reqwest::Client, acc: &Account, text: &str) -> Option<String> {
    let r: Value = c
        .post(format!("{PDS_URL}/xrpc/com.atproto.repo.createRecord"))
        .bearer_auth(&acc.access_jwt)
        .json(&json!({
            "repo": acc.did,
            "collection": "app.twi.post",
            "record": { "text": text, "createdAt": "2026-09-10T12:00:00Z" },
        }))
        .send()
        .await
        .ok()?
        .json()
        .await
        .ok()?;
    r["uri"].as_str().map(str::to_string)
}

type Ws = tokio_tungstenite::WebSocketStream<
    tokio_tungstenite::MaybeTlsStream<tokio::net::TcpStream>,
>;

/// Subscribe to the PDS firehose, or `None` if the endpoint isn't there
/// yet (see the module docs).
async fn subscribe() -> Option<Ws> {
    let url = pds_firehose::subscribe_url(PDS_WS, None);
    match tokio::time::timeout(Duration::from_secs(5), tokio_tungstenite::connect_async(&url)).await
    {
        Ok(Ok((ws, _))) => Some(ws),
        Ok(Err(e)) => {
            eprintln!(
                "cannot subscribe to {url}: {e} — the PDS endpoint com.atproto.sync.\
                 subscribeRepos is probably not implemented yet; skipping"
            );
            None
        }
        Err(_) => {
            eprintln!("timed out connecting to {url}; skipping");
            None
        }
    }
}

/// Read frames until one is a `#commit` for `did`, or the deadline
/// passes. `#info` frames along the way are tolerated (a fresh
/// subscription may legitimately be told its cursor is outdated).
async fn next_commit_for(
    ws: &mut Ws,
    did: &str,
    timeout: Duration,
) -> Option<pds_firehose::CommitFrame> {
    let deadline = tokio::time::Instant::now() + timeout;
    loop {
        let remaining = deadline.saturating_duration_since(tokio::time::Instant::now());
        if remaining.is_zero() {
            return None;
        }
        let msg = match tokio::time::timeout(remaining, ws.next()).await {
            Ok(Some(Ok(m))) => m,
            Ok(Some(Err(e))) => {
                eprintln!("firehose read error: {e}");
                return None;
            }
            Ok(None) | Err(_) => return None,
        };
        let Message::Binary(bytes) = msg else { continue };
        match pds_firehose::decode_frame(&bytes) {
            Ok(Frame::Commit(commit)) if commit.repo == did => return Some(*commit),
            Ok(Frame::Commit(_)) => continue,
            Ok(Frame::Info { name, .. }) => {
                eprintln!("firehose #info: {name}");
                continue;
            }
            Ok(Frame::Error { error, message }) => {
                eprintln!("firehose error frame: {error} {message:?}");
                return None;
            }
            Ok(Frame::Other { .. }) => continue,
            Err(e) => {
                // A frame we cannot decode is a contract failure worth
                // failing the test over — but only once we know the
                // endpoint exists, which we do by this point.
                panic!("could not decode a real PDS firehose frame: {e:#}");
            }
        }
    }
}

/// Everything a live test needs, or `None` with a printed reason.
async fn ready() -> Option<(reqwest::Client, PgPool, Account, Ws)> {
    let c = client();
    if !service_up(&c, PDS_URL).await {
        eprintln!("pds not running on {PDS_URL}, skipping");
        return None;
    }
    let Some(db) = appview_db().await else {
        eprintln!("DATABASE_URL_APPVIEW unset or unreachable, skipping");
        return None;
    };
    let Some(acc) = create_account(&c).await else {
        eprintln!("could not create a PDS account, skipping");
        return None;
    };
    // Subscribe *before* writing anything, so the commit we are about
    // to make is guaranteed to fall inside the subscription window.
    let ws = subscribe().await?;
    Some((c, db, acc, ws))
}

#[tokio::test]
async fn frame_from_the_local_pds_indexes_a_post() {
    let Some((c, db, acc, mut ws)) = ready().await else {
        return;
    };
    let text = format!("firehose e2e {}", uuid::Uuid::new_v4().simple());
    let Some(uri) = create_post(&c, &acc, &text).await else {
        eprintln!("createRecord failed, skipping");
        return;
    };

    let Some(commit) = next_commit_for(&mut ws, &acc.did, Duration::from_secs(15)).await else {
        eprintln!("no #commit frame for {} arrived in time, skipping", acc.did);
        return;
    };

    // The frame itself must carry what the contract promises.
    assert!(commit.seq > 0, "seq must be a positive sequence number");
    assert_eq!(commit.repo, acc.did);
    assert!(!commit.rev.is_empty(), "commit frames carry a rev");
    assert!(
        !commit.blocks.is_empty(),
        "a create commit must inline its record block"
    );
    let create = commit
        .ops
        .iter()
        .find(|op| op.action == "create" && op.collection() == Some("app.twi.post"))
        .expect("the frame must contain the post create op");
    assert!(create.cid.is_some(), "a create op carries the record CID");

    // The blocks field must be a CAR our reader understands, and the
    // op's CID must resolve inside it.
    let car = appview::car::parse(&commit.blocks).expect("blocks must be a readable CAR v1");
    assert!(
        car.block_map().contains_key(&create.cid.unwrap()),
        "the record block must be present in the CAR"
    );

    // And the whole path — frame → CAR → indexer — must land the row.
    let events = pds_firehose::events_from_frame(&commit).expect("events");
    assert!(
        events
            .iter()
            .any(|e| e.commit.as_ref().unwrap()["record"]["text"] == json!(text)),
        "the decoded record must carry the text we posted"
    );

    // Remove whatever the AppView's own push path already wrote, so the
    // assertion below is about *this* code applying *this* frame.
    sqlx::query("DELETE FROM posts WHERE uri = $1")
        .bind(&uri)
        .execute(&db)
        .await
        .unwrap();

    pds_firehose::apply_frame(&db, &commit)
        .await
        .expect("apply_frame");

    let stored: Option<String> = sqlx::query_scalar("SELECT text FROM posts WHERE uri = $1")
        .bind(&uri)
        .fetch_optional(&db)
        .await
        .unwrap();
    assert_eq!(
        stored.as_deref(),
        Some(text.as_str()),
        "the firehose frame must index the post at {uri}"
    );

    sqlx::query("DELETE FROM posts WHERE uri = $1")
        .bind(&uri)
        .execute(&db)
        .await
        .unwrap();
}

#[tokio::test]
async fn replaying_the_same_frame_changes_nothing() {
    let Some((c, db, acc, mut ws)) = ready().await else {
        return;
    };
    let text = format!("firehose replay {}", uuid::Uuid::new_v4().simple());
    let Some(uri) = create_post(&c, &acc, &text).await else {
        eprintln!("createRecord failed, skipping");
        return;
    };
    let Some(commit) = next_commit_for(&mut ws, &acc.did, Duration::from_secs(15)).await else {
        eprintln!("no #commit frame for {} arrived in time, skipping", acc.did);
        return;
    };

    // Three applications: the push already ran, then the firehose, then
    // a post-restart replay of the same seq.
    for _ in 0..3 {
        pds_firehose::apply_frame(&db, &commit)
            .await
            .expect("apply_frame");
    }

    let rows: i64 = sqlx::query_scalar("SELECT count(*) FROM posts WHERE uri = $1")
        .bind(&uri)
        .fetch_one(&db)
        .await
        .unwrap();
    assert_eq!(rows, 1, "replay must not duplicate {uri}");

    sqlx::query("DELETE FROM posts WHERE uri = $1")
        .bind(&uri)
        .execute(&db)
        .await
        .unwrap();
}

#[tokio::test]
async fn car_reader_parses_a_real_repo_export() {
    // This one needs no firehose: `getRepo` has always served a CAR
    // produced by the PDS's own writer, which is exactly the encoder
    // the firehose's `blocks` field reuses.
    let c = client();
    if !service_up(&c, PDS_URL).await {
        eprintln!("pds not running on {PDS_URL}, skipping");
        return;
    }
    let Some(acc) = create_account(&c).await else {
        eprintln!("could not create a PDS account, skipping");
        return;
    };
    if create_post(&c, &acc, "car reader fixture").await.is_none() {
        eprintln!("createRecord failed, skipping");
        return;
    }

    let resp = c
        .get(format!("{PDS_URL}/xrpc/com.atproto.sync.getRepo"))
        .query(&[("did", acc.did.as_str())])
        .send()
        .await
        .unwrap();
    if !resp.status().is_success() {
        eprintln!("getRepo returned {}, skipping", resp.status());
        return;
    }
    let bytes = resp.bytes().await.unwrap();
    let car = appview::car::parse(&bytes).expect("getRepo must return a readable CAR v1");
    assert_eq!(car.header.version, 1);
    assert!(
        !car.blocks.is_empty(),
        "a repo with one record has blocks (commit + MST + record)"
    );
    // Every block must hash to the CID the file declares — the strongest
    // available statement that the reader's section framing is right.
    appview::car::verify_block_cids(&car).expect("block CIDs must verify");
}

#[tokio::test]
async fn healthz_reports_the_firehose_state() {
    let c = client();
    if !service_up(&c, &appview_url()).await {
        eprintln!("appview not running, skipping");
        return;
    }
    let body: Value = c
        .get(format!("{}/healthz", appview_url()))
        .send()
        .await
        .unwrap()
        .json()
        .await
        .unwrap();

    // A binary built before this feature has none of these keys; say so
    // rather than failing, because "restart the AppView" is the fix.
    let Some(enabled) = body.get("pds_firehose_enabled").and_then(Value::as_bool) else {
        eprintln!(
            "the running AppView predates the PDS firehose (no pds_firehose_enabled \
             in /healthz) — rebuild and restart it; skipping"
        );
        return;
    };
    assert!(
        body.get("pds_firehose_connected")
            .and_then(Value::as_bool)
            .is_some(),
        "/healthz must report pds_firehose_connected: {body}"
    );
    assert!(
        body.get("pds_firehose_seq").and_then(Value::as_i64).is_some(),
        "/healthz must report pds_firehose_seq: {body}"
    );

    if !enabled {
        eprintln!("PDS_FIREHOSE_ENABLED=false on the running AppView; nothing more to check");
        return;
    }
    if !body["pds_firehose_connected"].as_bool().unwrap() {
        eprintln!(
            "the AppView is not connected to the PDS firehose — expected while \
             com.atproto.sync.subscribeRepos is still being implemented; skipping \
             the live-consumption check"
        );
        return;
    }

    // Connected: a new record must move the sequence number the AppView
    // reports, which is the end-to-end proof that the *service* (not
    // just this test process) consumes the stream.
    if !service_up(&c, PDS_URL).await {
        eprintln!("pds not running, skipping the live-consumption check");
        return;
    }
    let Some(acc) = create_account(&c).await else {
        eprintln!("could not create a PDS account, skipping");
        return;
    };
    let before = body["pds_firehose_seq"].as_i64().unwrap_or(0);
    if create_post(&c, &acc, "healthz seq probe").await.is_none() {
        eprintln!("createRecord failed, skipping");
        return;
    }
    for _ in 0..40 {
        tokio::time::sleep(Duration::from_millis(250)).await;
        let now: Value = c
            .get(format!("{}/healthz", appview_url()))
            .send()
            .await
            .unwrap()
            .json()
            .await
            .unwrap();
        if now["pds_firehose_seq"].as_i64().unwrap_or(0) > before {
            return; // consumed
        }
    }
    panic!("the AppView reports the firehose connected but its seq never advanced");
}