use std::time::Duration;
use headgate_core::{AdmitRequest, Envelope, Store, Transactional};
use headgate_postgres::PgStore;
fn conninfo() -> Option<String> {
match std::env::var("HG_TEST_PG") {
Ok(s) => Some(s),
Err(_) => {
eprintln!("HG_TEST_PG not set; skipping ORM interop matrix (postgres)");
None
}
}
}
fn scope() -> String {
let nanos = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.map(|d| d.as_nanos())
.unwrap_or(0);
format!("{}_{nanos}", std::process::id())
}
async fn app_client(conninfo: &str) -> tokio_postgres::Client {
let (client, conn) = tokio_postgres::connect(conninfo, tokio_postgres::NoTls)
.await
.expect("app connection");
tokio::spawn(async move {
let _ = conn.await;
});
client
}
fn env(queue: &str, id: &str) -> Envelope {
Envelope {
id: id.into(),
kind: "orm:t".into(),
payload: b"{}".to_vec(),
queue: queue.into(),
fingerprint: headgate_core::fingerprint("orm:t", b"{}"),
scheduled_at_ms: 1,
retention_ms: 86_400_000,
..Default::default()
}
}
async fn clean(c: &tokio_postgres::Client, queue: &str, app_table: &str) {
let _ = c
.batch_execute(&format!(
"DROP TABLE IF EXISTS {app_table};
DELETE FROM headgate_job WHERE queue = '{queue}';
DELETE FROM headgate_active_partition WHERE queue = '{queue}';
DELETE FROM headgate_effect WHERE key LIKE '{queue}-%'"
))
.await;
}
fn admit_req(queue: &str, worker: &str, lease: &str) -> AdmitRequest {
AdmitRequest {
worker: worker.into(),
lease_id: lease.into(),
queues: vec![queue.into()],
capacity: 10,
lease: Duration::from_secs(60),
quantum: 1000,
}
}
#[tokio::test]
async fn caller_tx_commit_makes_the_app_row_and_the_job_visible_and_admittable() {
let Some(conninfo) = conninfo() else { return };
let store = PgStore::connect(&conninfo, 2).expect("connect");
let s = scope();
let queue = format!("ormpg-a-{s}");
let app = format!("hg_orm_app_a_{s}");
let mut client = app_client(&conninfo).await;
clean(&client, &queue, &app).await;
client
.batch_execute(&format!(
"CREATE TABLE {app} (id text primary key, note text)"
))
.await
.expect("create app table");
let tx = client.transaction().await.expect("begin (caller-owned)");
tx.execute(
&format!("INSERT INTO {app} (id, note) VALUES ($1, $2)"),
&[&"order-1", &"paid"],
)
.await
.expect("app write");
store
.enqueue_on(&tx, &[env(&queue, &format!("{queue}-j1"))])
.await
.expect("enqueue_on the caller's tokio_postgres::Transaction");
tx.commit().await.expect("commit");
let app_rows: i64 = client
.query_one(&format!("SELECT count(*)::bigint FROM {app}"), &[])
.await
.expect("count app")
.get(0);
assert_eq!(app_rows, 1, "the app write must survive the commit");
let jobs: i64 = client
.query_one(
"SELECT count(*)::bigint FROM headgate_job WHERE queue = $1",
&[&queue],
)
.await
.expect("count jobs")
.get(0);
assert_eq!(jobs, 1, "the enqueue must survive the same commit");
let units = store
.admit(admit_req(&queue, "orm-w", "ORM-L1"))
.await
.expect("admit");
let ids: Vec<&str> = units
.iter()
.flat_map(|u| &u.claims)
.map(|c| c.envelope.id.as_str())
.collect();
assert_eq!(
ids,
vec![format!("{queue}-j1")],
"the committed job must be admittable"
);
clean(&client, &queue, &app).await;
}
#[tokio::test]
async fn caller_tx_rollback_leaves_neither_the_app_row_nor_the_job() {
let Some(conninfo) = conninfo() else { return };
let store = PgStore::connect(&conninfo, 2).expect("connect");
let s = scope();
let queue = format!("ormpg-b-{s}");
let app = format!("hg_orm_app_b_{s}");
let mut client = app_client(&conninfo).await;
clean(&client, &queue, &app).await;
client
.batch_execute(&format!(
"CREATE TABLE {app} (id text primary key, note text)"
))
.await
.expect("create app table");
let tx = client.transaction().await.expect("begin (caller-owned)");
tx.execute(
&format!("INSERT INTO {app} (id, note) VALUES ($1, $2)"),
&[&"order-2", &"pending"],
)
.await
.expect("app write");
store
.enqueue_on(&tx, &[env(&queue, &format!("{queue}-j1"))])
.await
.expect("enqueue_on the caller's transaction");
tx.rollback().await.expect("rollback");
let app_rows: i64 = client
.query_one(&format!("SELECT count(*)::bigint FROM {app}"), &[])
.await
.expect("count app")
.get(0);
assert_eq!(app_rows, 0, "the app write must be gone");
let jobs: i64 = client
.query_one(
"SELECT count(*)::bigint FROM headgate_job WHERE queue = $1",
&[&queue],
)
.await
.expect("count jobs")
.get(0);
assert_eq!(jobs, 0, "the enqueue must be gone WITH it — neither exists");
let units = store
.admit(admit_req(&queue, "orm-w", "ORM-L2"))
.await
.expect("admit");
assert!(
units.iter().flat_map(|u| &u.claims).next().is_none(),
"a rolled-back enqueue must never be admittable"
);
clean(&client, &queue, &app).await;
}
async fn deliver(
store: &PgStore,
txs: &dyn Transactional,
lease: &headgate_core::LeaseRef,
key: &str,
app: &str,
) -> bool {
let mut tx = store.begin().await.expect("begin");
let claimed = txs.claim_effect(&mut tx, key).await.expect("claim_effect");
if !claimed {
tx.rollback().await.expect("rollback");
return false; }
tx.client()
.expect("client")
.execute(
&format!("INSERT INTO {app} (id, note) VALUES ($1, $2)"),
&[&"charge-1", &"applied"],
)
.await
.expect("app write");
txs.complete_tx(&mut tx, lease).await.expect("complete_tx");
tx.commit().await.expect("commit");
true
}
#[tokio::test]
async fn once_in_a_caller_tx_does_not_double_apply_after_a_crash() {
let Some(conninfo) = conninfo() else { return };
let store = PgStore::connect(&conninfo, 2).expect("connect");
let txs = store.as_transactional().expect("postgres is transactional");
let s = scope();
let queue = format!("ormpg-c-{s}");
let app = format!("hg_orm_app_c_{s}");
let client = app_client(&conninfo).await;
clean(&client, &queue, &app).await;
client
.batch_execute(&format!(
"CREATE TABLE {app} (id text primary key, note text)"
))
.await
.expect("create app table");
let job_id = format!("{queue}-j1");
store
.enqueue(&[env(&queue, &job_id)])
.await
.expect("enqueue");
let units = store
.admit(admit_req(&queue, "orm-w", "ORM-L3"))
.await
.expect("admit");
let lease = units[0].claims[0].lease_ref();
let effect_key = format!("{queue}-effect");
assert!(
deliver(&store, txs, &lease, &effect_key, &app).await,
"first delivery runs the effect"
);
assert!(
!deliver(&store, txs, &lease, &effect_key, &app).await,
"a redelivery after a committed effect must skip the work entirely"
);
let charges: i64 = client
.query_one(&format!("SELECT count(*)::bigint FROM {app}"), &[])
.await
.expect("count app")
.get(0);
assert_eq!(charges, 1, "the app effect must be applied EXACTLY once");
let effects: i64 = client
.query_one(
"SELECT count(*)::bigint FROM headgate_effect WHERE key = $1",
&[&effect_key],
)
.await
.expect("count effects")
.get(0);
assert_eq!(effects, 1, "one effect-key row, claimed once, forever");
let state: String = client
.query_one(
"SELECT state::text FROM headgate_job WHERE ulid = $1",
&[&job_id],
)
.await
.expect("job row")
.get(0);
assert_eq!(
state, "completed",
"completion committed with the app write, not after it"
);
clean(&client, &queue, &app).await;
}