Skip to main content

ares_test_support/
lib.rs

1//! Shared live-test database harness for ARES integration tests.
2//!
3//! One resolver, one lifecycle, one failure policy for every crate's live
4//! Postgres tests (`ares_test`):
5//!
6//! - URL resolution: `TEST_DATABASE_URL` -> `DATABASE_URL` (postgres URLs only,
7//!   rewritten to `ares_test`) -> unix-socket peer auth, no credentials needed.
8//! - Lifecycle: once per test binary — connect, run migrations, truncate; the
9//!   tables are truncated again when the binary exits, so `ares_test` is left
10//!   empty between runs.
11//! - Failure policy: a configured-but-unreachable database fails loudly with
12//!   fix instructions instead of panicking mid-test or silently skipping.
13
14use std::future::Future;
15use std::sync::atomic::{AtomicBool, Ordering};
16
17static INIT: tokio::sync::OnceCell<()> = tokio::sync::OnceCell::const_new();
18static READY: AtomicBool = AtomicBool::new(false);
19
20/// Application tables truncated before and after each test binary run.
21const CLEANUP_TABLES: &[&str] = &[
22    "messages",
23    "conversations",
24    "sessions",
25    "memory_facts",
26    "preferences",
27    "user_agents",
28    "users",
29];
30
31/// Load `.env` once per test process, so `DATABASE_URL` is available if present.
32fn ensure_env_loaded() {
33    static ONCE: std::sync::Once = std::sync::Once::new();
34    ONCE.call_once(|| {
35        let _ = dotenvy::dotenv();
36    });
37}
38
39/// Returns the test database URL.
40///
41/// Priority:
42/// 1. `TEST_DATABASE_URL` env var (CI / custom setups)
43/// 2. `DATABASE_URL` env var rewritten to `ares_test` (postgres URLs only;
44///    non-postgres values such as the sqlite path in `.env.example` are ignored)
45/// 3. Fallback: unix-socket peer auth (`/var/run/postgresql`), which works with
46///    zero configuration for the OS user that owns `ares_test`
47pub fn test_db_url() -> String {
48    ensure_env_loaded();
49
50    if let Ok(url) = std::env::var("TEST_DATABASE_URL") {
51        return url;
52    }
53    if let Ok(url) = std::env::var("DATABASE_URL") {
54        if url.starts_with("postgres") {
55            if url.contains("/ares") && !url.contains("ares_test") {
56                return url.replace("/ares", "/ares_test");
57            }
58            return url;
59        }
60    }
61    let user = std::env::var("USER").unwrap_or_else(|_| "postgres".into());
62    format!("postgres://{user}@%2Fvar%2Frun%2Fpostgresql/ares_test")
63}
64
65/// Panic with actionable fix instructions for an unreachable test database.
66fn unreachable_panic(url: &str, reason: impl std::fmt::Display) -> ! {
67    panic!(
68        "live test DB unreachable at {url}\n\
69         Fix one of:\n  \
70         1. start postgres: sudo systemctl start postgresql\n  \
71         2. create the test DB: sudo -u postgres psql -c \"CREATE DATABASE ares_test OWNER $USER;\"\n  \
72         3. export TEST_DATABASE_URL=postgres://user:pass@host/ares_test\n\
73         underlying error: {reason}"
74    )
75}
76
77/// Truncate all application tables.
78async fn truncate_all(pool: &sqlx::PgPool) {
79    for table in CLEANUP_TABLES {
80        let query = format!("TRUNCATE TABLE {table} CASCADE");
81        if let Err(e) = sqlx::query(&query).execute(pool).await {
82            eprintln!("Warning: failed to truncate {table}: {e}");
83        }
84    }
85}
86
87fn connect_or_panic(url: &str) -> impl Future<Output = sqlx::PgPool> {
88    let url = url.to_string();
89    async move {
90        sqlx::postgres::PgPoolOptions::new()
91            .max_connections(5)
92            .connect(&url)
93            .await
94            .unwrap_or_else(|e| unreachable_panic(&url, e))
95    }
96}
97
98async fn init() {
99    let url = test_db_url();
100    let pool = connect_or_panic(&url).await;
101    if let Err(e) = ares_store::MIGRATOR.run(&pool).await {
102        unreachable_panic(&url, format!("migration failed: {e}"));
103    }
104    truncate_all(&pool).await;
105    READY.store(true, Ordering::SeqCst);
106}
107
108/// Connect to `ares_test`.
109///
110/// First call per test binary: connect, run migrations, truncate stale data.
111/// Every call returns a pool owned by the caller's runtime — test binaries run
112/// many short-lived tokio runtimes, and a pool shared across them outlives the
113/// runtime it was created on. An unreachable database panics with fix
114/// instructions — live coverage is never silently skipped.
115pub async fn pool() -> sqlx::PgPool {
116    INIT.get_or_init(init).await;
117    connect_or_panic(&test_db_url()).await
118}
119
120/// A fresh connection wrapped as [`ares_store::PostgresClient`].
121pub async fn client() -> ares_store::PostgresClient {
122    ares_store::PostgresClient { pool: pool().await }
123}
124
125/// Truncate again when the test binary exits, leaving `ares_test` empty.
126/// Runs only if this binary actually connected; teardown errors are reported
127/// but never mask test results.
128#[dtor::dtor]
129unsafe fn truncate_after_run() {
130    if !READY.load(Ordering::SeqCst) {
131        return;
132    }
133    let runtime = match tokio::runtime::Builder::new_current_thread()
134        .enable_all()
135        .build()
136    {
137        Ok(rt) => rt,
138        Err(e) => {
139            eprintln!("Warning: test-support teardown could not build a runtime: {e}");
140            return;
141        }
142    };
143    // A fresh connection is required here: connections held by `POOL` belong to
144    // the (now finished) test runtime and cannot be reused after main returns.
145    runtime.block_on(async {
146        let pool = match sqlx::postgres::PgPoolOptions::new()
147            .max_connections(1)
148            .acquire_timeout(std::time::Duration::from_secs(10))
149            .connect(&test_db_url())
150            .await
151        {
152            Ok(pool) => pool,
153            Err(e) => {
154                eprintln!("Warning: test-support teardown could not connect: {e}");
155                return;
156            }
157        };
158        truncate_all(&pool).await;
159        pool.close().await;
160    });
161}
162
163#[cfg(test)]
164mod tests {
165    use super::*;
166
167    fn env_lock() -> std::sync::MutexGuard<'static, ()> {
168        static LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(());
169        LOCK.lock().unwrap_or_else(|p| p.into_inner())
170    }
171
172    #[test]
173    fn fallback_is_socket_peer_auth() {
174        let _g = env_lock();
175        let saved_test = std::env::var("TEST_DATABASE_URL").ok();
176        let saved_db = std::env::var("DATABASE_URL").ok();
177        std::env::remove_var("TEST_DATABASE_URL");
178        std::env::remove_var("DATABASE_URL");
179        let url = test_db_url();
180        if let Some(v) = saved_test {
181            std::env::set_var("TEST_DATABASE_URL", v);
182        }
183        if let Some(v) = saved_db {
184            std::env::set_var("DATABASE_URL", v);
185        }
186        assert!(url.contains("%2Fvar%2Frun%2Fpostgresql"), "unexpected: {url}");
187        assert!(url.ends_with("/ares_test"), "unexpected: {url}");
188        assert!(!url.contains("localhost"), "no TCP fallback expected: {url}");
189    }
190
191    #[test]
192    fn sqlite_database_url_is_ignored() {
193        let _g = env_lock();
194        let saved_test = std::env::var("TEST_DATABASE_URL").ok();
195        let saved_db = std::env::var("DATABASE_URL").ok();
196        std::env::remove_var("TEST_DATABASE_URL");
197        std::env::set_var("DATABASE_URL", "./data/ares.db");
198        let url = test_db_url();
199        if let Some(v) = saved_test {
200            std::env::set_var("TEST_DATABASE_URL", v);
201        }
202        if let Some(v) = saved_db {
203            std::env::set_var("DATABASE_URL", v);
204        }
205        assert!(url.contains("%2Fvar%2Frun%2Fpostgresql"), "unexpected: {url}");
206    }
207
208    /// Local smoke test: run with `cargo test -p ares-test-support -- --ignored`.
209    #[tokio::test]
210    #[ignore = "requires a live ares_test database"]
211    async fn connects_via_socket_without_credentials() {
212        std::env::remove_var("TEST_DATABASE_URL");
213        std::env::remove_var("DATABASE_URL");
214        let pool = pool().await;
215        let row: (i32,) = sqlx::query_as("SELECT 1").fetch_one(&pool).await.unwrap();
216        assert_eq!(row.0, 1);
217    }
218}