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/// Clear any previously-activated runtime backend so legacy integration tests
53/// start from the canonical compatibility surface unless they opt into a
54/// queue-storage runtime explicitly.
55pub async fn reset_runtime_backend(pool: &PgPool) {
56    let mut tx = pool
57        .begin()
58        .await
59        .expect("Failed to start runtime backend reset transaction");
60
61    sqlx::query(
62        r#"
63        UPDATE awa.storage_transition_state
64        SET current_engine = 'canonical',
65            prepared_engine = NULL,
66            state = 'canonical',
67            transition_epoch = transition_epoch + 1,
68            details = '{}'::jsonb,
69            updated_at = now(),
70            finalized_at = NULL
71        WHERE singleton
72        "#,
73    )
74    .execute(&mut *tx)
75    .await
76    .expect("Failed to reset storage transition state for test setup");
77    sqlx::query("DELETE FROM awa.runtime_storage_backends WHERE backend = 'queue_storage'")
78        .execute(&mut *tx)
79        .await
80        .expect("Failed to reset active runtime backend for test setup");
81    sqlx::query("DELETE FROM awa.runtime_instances")
82        .execute(&mut *tx)
83        .await
84        .expect("Failed to reset runtime instances for test setup");
85    tx.commit()
86        .await
87        .expect("Failed to commit runtime backend reset transaction");
88}
89
90/// Delete all jobs, queue metadata, and admin caches for a specific queue.
91///
92/// Explicitly deletes the `queue_state_counts` row to prevent accumulated
93/// cache drift from affecting assertions. The DELETE trigger normally handles
94/// this, but concurrent test runs against a shared DB can cause small delta
95/// errors that compound over time.
96pub async fn clean_queue(pool: &PgPool, queue: &str) {
97    sqlx::query("DELETE FROM awa.jobs WHERE queue = $1")
98        .bind(queue)
99        .execute(pool)
100        .await
101        .expect("Failed to clean queue jobs");
102    sqlx::query("DELETE FROM awa.queue_meta WHERE queue = $1")
103        .bind(queue)
104        .execute(pool)
105        .await
106        .expect("Failed to clean queue meta");
107    sqlx::query("DELETE FROM awa.queue_state_counts WHERE queue = $1")
108        .bind(queue)
109        .execute(pool)
110        .await
111        .expect("Failed to clean queue state counts");
112}
113
114/// Query job state counts for a queue, returning a map of state -> count.
115pub async fn queue_state_counts(pool: &PgPool, queue: &str) -> HashMap<String, i64> {
116    let rows: Vec<(String, i64)> = sqlx::query_as(
117        r#"
118        SELECT state::text, count(*)::bigint
119        FROM awa.jobs
120        WHERE queue = $1
121        GROUP BY state
122        "#,
123    )
124    .bind(queue)
125    .fetch_all(pool)
126    .await
127    .expect("Failed to query state counts");
128
129    rows.into_iter().collect()
130}
131
132/// Extract a count for a given state from a state-counts map.
133pub fn state_count(counts: &HashMap<String, i64>, state: &str) -> i64 {
134    counts.get(state).copied().unwrap_or(0)
135}
136
137/// Poll queue state counts until a predicate is satisfied, or panic on timeout.
138pub async fn wait_for_counts(
139    pool: &PgPool,
140    queue: &str,
141    predicate: impl Fn(&HashMap<String, i64>) -> bool,
142    timeout: Duration,
143) -> HashMap<String, i64> {
144    let start = std::time::Instant::now();
145    loop {
146        let counts = queue_state_counts(pool, queue).await;
147        if predicate(&counts) {
148            return counts;
149        }
150        assert!(
151            start.elapsed() < timeout,
152            "Timed out waiting for queue {queue} counts; last counts: {counts:?}"
153        );
154        tokio::time::sleep(Duration::from_millis(50)).await;
155    }
156}
awa_testing/setup.rs

awa_testing/
setup.rs
