github-ops
134 lines
3.7 KiB
Rust
134 lines
3.7 KiB
Rust
use std::{
|
|
process::{Command, Stdio},
|
|
sync::{Mutex, OnceLock},
|
|
};
|
|
|
|
use sqlx::{Connection, Executor, PgConnection};
|
|
use testcontainers::{Container, ImageExt, runners::SyncRunner};
|
|
use testcontainers_modules::postgres::Postgres;
|
|
use tokio::sync::OnceCell;
|
|
use tokio::time::{Duration, sleep};
|
|
|
|
static POSTGRES: OnceCell<TestPostgres> = OnceCell::const_new();
|
|
static CLEANUP_REGISTERED: OnceLock<()> = OnceLock::new();
|
|
static CONTAINER_IDS: Mutex<Vec<String>> = Mutex::new(Vec::new());
|
|
|
|
struct TestPostgres {
|
|
_container: Container<Postgres>,
|
|
database_url: String,
|
|
}
|
|
|
|
pub async fn postgres_database_url() -> &'static str {
|
|
&test_postgres().await.database_url
|
|
}
|
|
|
|
pub async fn postgres_schema_url(prefix: &str) -> String {
|
|
let database_url = postgres_database_url().await;
|
|
let schema = unique_schema_name(prefix);
|
|
let mut connection = connect_admin_with_retry(database_url).await;
|
|
|
|
connection
|
|
.execute(sqlx::query(&format!("create schema {schema}")))
|
|
.await
|
|
.expect("test PostgreSQL schema must be created");
|
|
|
|
format!("{database_url}?options=-csearch_path%3D{schema}")
|
|
}
|
|
|
|
async fn connect_admin_with_retry(database_url: &str) -> PgConnection {
|
|
let mut last_error = None;
|
|
|
|
for _ in 0..80 {
|
|
match PgConnection::connect(database_url).await {
|
|
Ok(connection) => return connection,
|
|
Err(error) => {
|
|
last_error = Some(error);
|
|
sleep(Duration::from_millis(250)).await;
|
|
}
|
|
}
|
|
}
|
|
|
|
panic!(
|
|
"test PostgreSQL container must accept admin connections after readiness wait: {}",
|
|
last_error
|
|
.map(|error| error.to_string())
|
|
.unwrap_or_else(|| "connection was not attempted".to_owned())
|
|
);
|
|
}
|
|
|
|
pub fn unique_schema_name(prefix: &str) -> String {
|
|
let normalized_prefix: String = prefix
|
|
.chars()
|
|
.map(|ch| {
|
|
if ch.is_ascii_alphanumeric() || ch == '_' {
|
|
ch.to_ascii_lowercase()
|
|
} else {
|
|
'_'
|
|
}
|
|
})
|
|
.collect();
|
|
|
|
format!("{}_{}", normalized_prefix, uuid::Uuid::now_v7().simple())
|
|
}
|
|
|
|
async fn test_postgres() -> &'static TestPostgres {
|
|
POSTGRES
|
|
.get_or_init(|| async {
|
|
tokio::task::spawn_blocking(start_postgres)
|
|
.await
|
|
.expect("test PostgreSQL container startup task must complete")
|
|
})
|
|
.await
|
|
}
|
|
|
|
fn start_postgres() -> TestPostgres {
|
|
let container = Postgres::default()
|
|
.with_db_name("crank")
|
|
.with_user("crank")
|
|
.with_password("crank")
|
|
.with_tag("16-alpine")
|
|
.start()
|
|
.expect("test PostgreSQL container must start");
|
|
register_container_cleanup(container.id().to_owned());
|
|
let host = container
|
|
.get_host()
|
|
.expect("test PostgreSQL host must be resolved");
|
|
let port = container
|
|
.get_host_port_ipv4(5432)
|
|
.expect("test PostgreSQL port must be mapped");
|
|
|
|
TestPostgres {
|
|
_container: container,
|
|
database_url: format!("postgres://crank:crank@{host}:{port}/crank"),
|
|
}
|
|
}
|
|
|
|
fn register_container_cleanup(container_id: String) {
|
|
CLEANUP_REGISTERED.get_or_init(|| unsafe {
|
|
libc::atexit(cleanup_test_containers);
|
|
});
|
|
CONTAINER_IDS
|
|
.lock()
|
|
.expect("test container id registry must be available")
|
|
.push(container_id);
|
|
}
|
|
|
|
extern "C" fn cleanup_test_containers() {
|
|
let container_ids = CONTAINER_IDS
|
|
.lock()
|
|
.map(|ids| ids.clone())
|
|
.unwrap_or_default();
|
|
|
|
if container_ids.is_empty() {
|
|
return;
|
|
}
|
|
|
|
let _ = Command::new("docker")
|
|
.arg("rm")
|
|
.arg("-f")
|
|
.args(container_ids)
|
|
.stdout(Stdio::null())
|
|
.stderr(Stdio::null())
|
|
.status();
|
|
}
|