1use 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
20const CLEANUP_TABLES: &[&str] = &[
22 "messages",
23 "conversations",
24 "sessions",
25 "memory_facts",
26 "preferences",
27 "user_agents",
28 "users",
29];
30
31fn ensure_env_loaded() {
33 static ONCE: std::sync::Once = std::sync::Once::new();
34 ONCE.call_once(|| {
35 let _ = dotenvy::dotenv();
36 });
37}
38
39pub 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
65fn 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
77async 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
108pub async fn pool() -> sqlx::PgPool {
116 INIT.get_or_init(init).await;
117 connect_or_panic(&test_db_url()).await
118}
119
120pub async fn client() -> ares_store::PostgresClient {
122 ares_store::PostgresClient { pool: pool().await }
123}
124
125#[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 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 #[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}