athena_rs 3.26.2

//! Test-only helpers exported for integration tests under `tests/`.
use std::sync::{Mutex, OnceLock};

static ATHENA_KEY_TEST_LOCK: OnceLock<Mutex<()>> = OnceLock::new();
static TEST_ENV_LOCK: OnceLock<Mutex<()>> = OnceLock::new();

/// `ATHENA_KEY_12` value used while [`AthAdminKeyGuard`] is held.
pub const ATHENA_TEST_ADMIN_KEY: &str = "integration-test-athena-admin-key-2026";

/// Serializes `ATHENA_KEY_12` across concurrent tests.
pub struct AthAdminKeyGuard {
    _guard: std::sync::MutexGuard<'static, ()>,
}

impl AthAdminKeyGuard {
    pub fn new() -> Self {
        let _guard = ATHENA_KEY_TEST_LOCK
            .get_or_init(|| Mutex::new(()))
            .lock()
            .unwrap_or_else(|err| err.into_inner());
        unsafe {
            std::env::set_var("ATHENA_KEY_12", ATHENA_TEST_ADMIN_KEY);
        }
        Self { _guard }
    }
}

impl Drop for AthAdminKeyGuard {
    fn drop(&mut self) {
        unsafe {
            std::env::remove_var("ATHENA_KEY_12");
        }
    }
}

/// Serializes arbitrary process env mutations across tests and restores the prior value on drop.
pub struct TestEnvGuard {
    _guard: std::sync::MutexGuard<'static, ()>,
    entries: Vec<(&'static str, Option<String>)>,
}

impl TestEnvGuard {
    pub fn set(key: &'static str, value: &str) -> Self {
        Self::apply(&[(key, Some(value))])
    }

    pub fn unset(key: &'static str) -> Self {
        Self::apply(&[(key, None)])
    }

    pub fn apply(updates: &[(&'static str, Option<&str>)]) -> Self {
        let guard = TEST_ENV_LOCK
            .get_or_init(|| Mutex::new(()))
            .lock()
            .unwrap_or_else(|err| err.into_inner());
        let mut entries: Vec<(&str, Option<String>)> = Vec::with_capacity(updates.len());
        for (key, value) in updates {
            entries.push((*key, std::env::var(key).ok()));
            match value {
                Some(value) => unsafe {
                    std::env::set_var(key, value);
                },
                None => unsafe {
                    std::env::remove_var(key);
                },
            }
        }
        Self {
            _guard: guard,
            entries,
        }
    }
}

impl Drop for TestEnvGuard {
    fn drop(&mut self) {
        for (key, previous) in self.entries.iter().rev() {
            match previous {
                Some(value) => unsafe {
                    std::env::set_var(key, value);
                },
                None => unsafe {
                    std::env::remove_var(key);
                },
            }
        }
    }
}