awa_testing/
setup.rs

1//! Common test setup utilities for Awa integration tests.
2
3use sqlx::postgres::PgPoolOptions;
4use sqlx::PgPool;
5use std::collections::HashMap;
6use std::time::Duration;
7
8/// Default database URL for test runs.
9pub fn database_url() -> String {
10    std::env::var("DATABASE_URL")
11        .unwrap_or_else(|_| "postgres://postgres:test@localhost:15432/awa_test".to_string())
12}
13
14/// Database URL with a custom application_name parameter appended.
15pub fn database_url_with_app_name(app_name: &str) -> String {
16    let mut url = database_url();
17    let sep = if url.contains('?') { '&' } else { '?' };
18    url.push(sep);
19    url.push_str("application_name=");
20    url.push_str(app_name);
21    url
22}
23
24/// Create a connection pool.
25pub async fn pool(max_connections: u32) -> PgPool {
26    PgPoolOptions::new()
27        .max_connections(max_connections)
28        .connect(&database_url())
29        .await
30        .expect("Failed to connect to database")
31}
32
33/// Create a connection pool with a custom database URL.
34pub async fn pool_with_url(url: &str, max_connections: u32) -> PgPool {
35    PgPoolOptions::new()
36        .max_connections(max_connections)
37        .connect(url)
38        .await
39        .expect("Failed to connect to database")
40}
41
42/// Create a pool, run migrations, and return it.
43pub async fn setup(max_connections: u32) -> PgPool {
44    let pool = pool(max_connections).await;
45    awa_model::migrations::run(&pool)
46        .await
47        .expect("Failed to run migrations");
48    reset_runtime_backend(&pool).await;
49    pool
50}
51
52/// Initialise the active storage engine to the one selected by the
53/// `AWA_TEST_ENGINE` env var: `canonical` (default) or `queue_storage`. awa's
54/// caller-facing contract is engine-invariant, so running the same suite under
55/// both values exercises both backends. Tests that pin a specific engine call
56/// [`reset_to_canonical`] / [`activate_queue_storage`] directly.
57pub async fn reset_runtime_backend(pool: &PgPool) {
58    match std::env::var("AWA_TEST_ENGINE").as_deref() {
59        Ok("queue_storage") => activate_queue_storage(pool).await,
60        _ => reset_to_canonical(pool).await,
61    }
62}
63
64/// Force the canonical engine and clear any queue-storage runtime registration.
65pub async fn reset_to_canonical(pool: &PgPool) {
66    let mut tx = pool
67        .begin()
68        .await
69        .expect("Failed to start runtime backend reset transaction");
70
71    sqlx::query(
72        r#"
73        UPDATE awa.storage_transition_state
74        SET current_engine = 'canonical',
75            prepared_engine = NULL,
76            state = 'canonical',
77            transition_epoch = transition_epoch + 1,
78            details = '{}'::jsonb,
79            updated_at = now(),
80            finalized_at = NULL
81        WHERE singleton
82        "#,
83    )
84    .execute(&mut *tx)
85    .await
86    .expect("Failed to reset storage transition state for test setup");
87    sqlx::query("DELETE FROM awa.runtime_storage_backends WHERE backend = 'queue_storage'")
88        .execute(&mut *tx)
89        .await
90        .expect("Failed to reset active runtime backend for test setup");
91    sqlx::query("DELETE FROM awa.runtime_instances")
92        .execute(&mut *tx)
93        .await
94        .expect("Failed to reset runtime instances for test setup");
95    tx.commit()
96        .await
97        .expect("Failed to commit runtime backend reset transaction");
98}
99
100/// Activate the queue-storage engine against the substrate installed in the
101/// `awa` schema (mirrors a fresh-install auto-finalize).
102pub async fn activate_queue_storage(pool: &PgPool) {
103    let mut tx = pool
104        .begin()
105        .await
106        .expect("Failed to start queue-storage activation transaction");
107
108    sqlx::query(
109        r#"
110        UPDATE awa.storage_transition_state
111        SET state = 'active',
112            current_engine = 'queue_storage',
113            prepared_engine = NULL,
114            details = jsonb_build_object('schema', 'awa'),
115            transition_epoch = transition_epoch + 1,
116            updated_at = now(),
117            finalized_at = now()
118        WHERE singleton
119        "#,
120    )
121    .execute(&mut *tx)
122    .await
123    .expect("Failed to activate queue-storage transition state for test setup");
124    sqlx::query(
125        r#"
126        INSERT INTO awa.runtime_storage_backends (backend, schema_name, updated_at)
127        VALUES ('queue_storage', 'awa', now())
128        ON CONFLICT (backend) DO UPDATE
129        SET schema_name = EXCLUDED.schema_name, updated_at = EXCLUDED.updated_at
130        "#,
131    )
132    .execute(&mut *tx)
133    .await
134    .expect("Failed to register queue-storage backend for test setup");
135    sqlx::query("DELETE FROM awa.runtime_instances")
136        .execute(&mut *tx)
137        .await
138        .expect("Failed to reset runtime instances for test setup");
139    tx.commit()
140        .await
141        .expect("Failed to commit queue-storage activation transaction");
142}
143
144/// The storage engine the suite is running against, selected by `AWA_TEST_ENGINE`.
145#[derive(Debug, Clone, Copy, PartialEq, Eq)]
146pub enum TestEngine {
147    Canonical,
148    QueueStorage,
149}
150
151/// The engine selected for this test run (`canonical` unless `AWA_TEST_ENGINE=queue_storage`).
152pub fn test_engine() -> TestEngine {
153    match std::env::var("AWA_TEST_ENGINE").as_deref() {
154        Ok("queue_storage") => TestEngine::QueueStorage,
155        _ => TestEngine::Canonical,
156    }
157}
158
159/// Guard for tests whose assertions are specific to the canonical engine —
160/// e.g. admin-metadata caches maintained by canonical row triggers, or the
161/// canonical running-cancel notification. Returns `true` (with a skip notice)
162/// when the run is under another engine, so the test can early-return:
163///
164/// ```ignore
165/// if awa_testing::setup::skip_unless_canonical("my_test") { return; }
166/// ```
167pub fn skip_unless_canonical(test: &str) -> bool {
168    if test_engine() != TestEngine::Canonical {
169        eprintln!(
170            "[skip] {test}: canonical-only assertions; running under AWA_TEST_ENGINE=queue_storage"
171        );
172        true
173    } else {
174        false
175    }
176}
177
178/// Delete all jobs, queue metadata, and admin caches for a specific queue.
179///
180/// Explicitly deletes the `queue_state_counts` row to prevent accumulated
181/// cache drift from affecting assertions. The DELETE trigger normally handles
182/// this, but concurrent test runs against a shared DB can cause small delta
183/// errors that compound over time.
184pub async fn clean_queue(pool: &PgPool, queue: &str) {
185    match awa_model::queue_storage::QueueStorage::active_schema(pool)
186        .await
187        .expect("active_schema lookup for clean_queue")
188    {
189        Some(schema) => clean_queue_substrate(pool, &schema, queue).await,
190        None => {
191            sqlx::query("DELETE FROM awa.jobs WHERE queue = $1")
192                .bind(queue)
193                .execute(pool)
194                .await
195                .expect("Failed to clean queue jobs");
196        }
197    }
198    sqlx::query("DELETE FROM awa.queue_meta WHERE queue = $1")
199        .bind(queue)
200        .execute(pool)
201        .await
202        .expect("Failed to clean queue meta");
203    sqlx::query("DELETE FROM awa.queue_state_counts WHERE queue = $1")
204        .bind(queue)
205        .execute(pool)
206        .await
207        .expect("Failed to clean queue state counts");
208}
209
210/// Drain every job-bearing plane of the queue-storage substrate for one queue
211/// and release its unique claims, so a queue name can be reused across tests
212/// and re-runs against the same database. The canonical `awa.jobs` view only
213/// surfaces a queue's head under queue storage, so a single view `DELETE`
214/// leaves the rest of the lane (and its `job_unique_claims`) behind.
215async fn clean_queue_substrate(pool: &PgPool, schema: &str, queue: &str) {
216    // Release unique claims first, while the rows that carry the job ids exist.
217    sqlx::query(&format!(
218        "DELETE FROM awa.job_unique_claims WHERE job_id IN (
219             SELECT job_id FROM {schema}.ready_entries WHERE queue = $1
220             UNION SELECT job_id FROM {schema}.deferred_jobs WHERE queue = $1
221             UNION SELECT job_id FROM {schema}.leases WHERE queue = $1
222             UNION SELECT job_id FROM {schema}.done_entries WHERE queue = $1
223             UNION SELECT job_id FROM {schema}.dlq_entries WHERE queue = $1)"
224    ))
225    .bind(queue)
226    .execute(pool)
227    .await
228    .expect("Failed to release queue unique claims");
229
230    for plane in [
231        "ready_entries",
232        "deferred_jobs",
233        "leases",
234        "done_entries",
235        "dlq_entries",
236    ] {
237        sqlx::query(&format!("DELETE FROM {schema}.{plane} WHERE queue = $1"))
238            .bind(queue)
239            .execute(pool)
240            .await
241            .unwrap_or_else(|err| panic!("Failed to clean {plane} for queue {queue}: {err}"));
242    }
243}
244
245/// Query job state counts for a queue, returning a map of state -> count.
246pub async fn queue_state_counts(pool: &PgPool, queue: &str) -> HashMap<String, i64> {
247    let rows: Vec<(String, i64)> = sqlx::query_as(
248        r#"
249        SELECT state::text, count(*)::bigint
250        FROM awa.jobs
251        WHERE queue = $1
252        GROUP BY state
253        "#,
254    )
255    .bind(queue)
256    .fetch_all(pool)
257    .await
258    .expect("Failed to query state counts");
259
260    rows.into_iter().collect()
261}
262
263/// Extract a count for a given state from a state-counts map.
264pub fn state_count(counts: &HashMap<String, i64>, state: &str) -> i64 {
265    counts.get(state).copied().unwrap_or(0)
266}
267
268/// Poll queue state counts until a predicate is satisfied, or panic on timeout.
269pub async fn wait_for_counts(
270    pool: &PgPool,
271    queue: &str,
272    predicate: impl Fn(&HashMap<String, i64>) -> bool,
273    timeout: Duration,
274) -> HashMap<String, i64> {
275    let start = std::time::Instant::now();
276    loop {
277        let counts = queue_state_counts(pool, queue).await;
278        if predicate(&counts) {
279            return counts;
280        }
281        assert!(
282            start.elapsed() < timeout,
283            "Timed out waiting for queue {queue} counts; last counts: {counts:?}"
284        );
285        tokio::time::sleep(Duration::from_millis(50)).await;
286    }
287}
awa_testing/setup.rs

awa_testing/
setup.rs
