martin/tests/utils/pg_utils.rs

pub use martin::args::Env;
use martin::pg::{PgConfig, Pool, TableInfo};
use martin::{IdResolver, Source, Sources};

use crate::FauxEnv;

//
// This file is used by many tests and benchmarks.
// Each function should allow dead_code as they might not be used by a specific test file.
//

pub type MockSource = (Sources, PgConfig);

#[allow(dead_code)]
#[must_use]
pub fn mock_pgcfg(yaml: &str) -> PgConfig {
    let Ok(db_url) = std::env::var("DATABASE_URL") else {
        panic!("DATABASE_URL env var is not set. Unable to do integration tests");
    };
    let env = FauxEnv(vec![("DATABASE_URL", db_url.into())].into_iter().collect());
    let mut cfg: PgConfig = subst::yaml::from_str(yaml, &env).unwrap();
    cfg.finalize().unwrap();
    cfg
}

#[allow(dead_code)]
pub async fn mock_pool() -> Pool {
    let cfg = mock_pgcfg("connection_string: $DATABASE_URL");
    let res = Pool::new(&cfg).await;
    res.expect("Failed to create pool")
}

#[allow(dead_code)]
pub async fn mock_sources(mut config: PgConfig) -> MockSource {
    let res = config.resolve(IdResolver::default()).await;
    let res = res.expect("Failed to resolve pg data");
    (res, config)
}

#[allow(dead_code)]
#[must_use]
pub fn table<'a>(mock: &'a MockSource, name: &str) -> &'a TableInfo {
    let (_, PgConfig { tables, .. }) = mock;
    tables.as_ref().map(|v| v.get(name).unwrap()).unwrap()
}

#[allow(dead_code)]
#[must_use]
pub fn source<'a>(mock: &'a MockSource, name: &str) -> &'a dyn Source {
    let (sources, _) = mock;
    sources.get(name).unwrap().as_ref()
}