pg-embed-setup-unpriv 0.5.2

//! Tests for environment scoping and logging.

use std::{
    env,
    ffi::{OsStr, OsString},
    panic,
    sync::{Arc, Barrier, TryLockError, mpsc},
    thread,
    time::{Duration, Instant},
};

use rstest::rstest;
use serial_test::serial;

use super::{
    ScopedEnv,
    THREAD_STATE,
    state::{ENV_LOCK, ThreadState},
};
#[cfg(feature = "cluster-unit-tests")]
use crate::test_support::capture_info_logs;

mod corruption;
mod thread_helpers;
mod validation;

use corruption::{
    CorruptionCase,
    apply_invalid_scope_exit,
    drop_guards_in_order,
    drop_guards_out_of_order,
    no_corruption,
    run_scoped_env_corruption_test,
    setup_nested_guards,
    setup_single_guard,
};
use thread_helpers::{
    ReleaseOnDrop,
    RestoreEnv,
    ThreadAChannels,
    ThreadBChannels,
    spawn_inner_guard_thread,
    spawn_outer_guard_thread,
};

#[test]
#[serial]
fn recovers_from_poisoned_lock() {
    assert!(
        panic::catch_unwind(|| {
            let _guard =
                ScopedEnv::apply(&[(String::from("POISON_TEST"), Some(String::from("one")))]);
            panic!("intentional panic to poison the mutex");
        })
        .is_err()
    );

    let guard = ScopedEnv::apply(&[(String::from("POISON_TEST"), Some(String::from("two")))]);
    assert_eq!(env::var("POISON_TEST").as_deref(), Ok("two"));
    drop(guard);
    assert!(env::var("POISON_TEST").is_err());
}

/// Verifies nested scopes restore the outer value after the inner drops.
#[test]
#[serial]
fn allows_reentrant_scopes() {
    let outer = ScopedEnv::apply(&[(String::from("NESTED_TEST"), Some(String::from("outer")))]);
    assert_eq!(env::var("NESTED_TEST").as_deref(), Ok("outer"));

    {
        let inner = ScopedEnv::apply(&[(String::from("NESTED_TEST"), Some(String::from("inner")))]);
        assert_eq!(env::var("NESTED_TEST").as_deref(), Ok("inner"));
        drop(inner);
    }

    assert_eq!(env::var("NESTED_TEST").as_deref(), Ok("outer"));
    drop(outer);
    assert!(env::var("NESTED_TEST").is_err());
}

/// Verifies the process lock remains held until the final nested scope exits.
#[test]
#[serial]
fn keeps_lock_until_last_scope_drops() {
    let outer = ScopedEnv::apply(&[(String::from("SCOPE_TEST"), Some(String::from("outer")))]);
    let inner = ScopedEnv::apply(&[(String::from("SCOPE_TEST"), Some(String::from("inner")))]);

    drop(outer);
    assert_eq!(env::var("SCOPE_TEST").as_deref(), Ok("inner"));
    assert!(
        ENV_LOCK.try_lock().is_err(),
        "mutex must remain held by inner guard"
    );

    let third = ScopedEnv::apply(&[(String::from("SCOPE_TEST"), Some(String::from("third")))]);
    assert_eq!(env::var("SCOPE_TEST").as_deref(), Ok("third"));
    drop(third);
    assert_eq!(env::var("SCOPE_TEST").as_deref(), Ok("inner"));

    drop(inner);
    let free = ENV_LOCK
        .try_lock()
        .expect("mutex should release after final scope drops");
    drop(free);
    assert!(env::var("SCOPE_TEST").is_err());
}

/// Verifies a second thread blocks until the first scoped guard releases.
#[test]
#[serial]
fn serializes_env_across_threads() {
    let key = "THREAD_SCOPE_TEST";
    let restore_env = RestoreEnv {
        key: String::from(key),
        original: env::var_os(key),
    };
    set_env_var_locked(OsStr::new(key), OsStr::new("pre-existing"));

    let (ready_tx, ready_rx) = mpsc::channel();
    let (start_tx, start_rx) = mpsc::channel();
    let (release_tx, release_rx) = mpsc::channel();
    let (attempt_tx, attempt_rx) = mpsc::channel();
    let (acquired_tx, acquired_rx) = mpsc::channel();
    let (done_a_tx, done_a_rx) = mpsc::channel();
    let (done_b_tx, done_b_rx) = mpsc::channel();
    let barrier = Arc::new(Barrier::new(2));
    // Generous timeout to avoid hanging tests, not to enforce ordering.
    let deadlock_timeout = Duration::from_secs(30);

    let thread_a = spawn_outer_guard_thread(
        String::from(key),
        ThreadAChannels {
            barrier: Arc::clone(&barrier),
            ready_tx,
            release_rx,
            done_tx: done_a_tx,
        },
    );
    let thread_b = spawn_inner_guard_thread(
        String::from(key),
        ThreadBChannels {
            start_rx,
            attempt_tx,
            acquired_tx,
            done_tx: done_b_tx,
        },
    );

    ready_rx
        .recv_timeout(deadlock_timeout)
        .expect("outer guard should be ready");
    barrier.wait();

    let release_guard = ReleaseOnDrop {
        sender: Some(release_tx),
    };

    start_tx.send(()).expect("start signal must be sent");
    attempt_rx
        .recv_timeout(deadlock_timeout)
        .expect("second thread should attempt to acquire the guard");

    assert!(
        acquired_rx.try_recv().is_err(),
        "second thread must block while the outer guard holds the lock"
    );

    drop(release_guard);

    let value = acquired_rx
        .recv_timeout(deadlock_timeout)
        .expect("second thread should acquire after release");
    assert_eq!(value.as_deref(), Some("two"));

    done_a_rx
        .recv_timeout(deadlock_timeout)
        .expect("outer guard thread should finish");
    done_b_rx
        .recv_timeout(deadlock_timeout)
        .expect("inner guard thread should finish");

    thread_a.join().expect("thread A should exit cleanly");
    thread_b.join().expect("thread B should exit cleanly");

    assert_eq!(env::var(key).as_deref(), Ok("pre-existing"));
    assert_env_lock_released();
    drop(restore_env);
}

/// Verifies invalid scope exits restore state without panicking.
#[test]
#[serial]
fn thread_state_recovers_from_invalid_index() {
    let key = OsString::from("THREAD_STATE_INVALID_INDEX");
    let original = env::var_os(&key);
    let mut state = ThreadState::new();
    let index = state.enter_scope(vec![(key.clone(), Some(OsString::from("value")))]);

    assert_eq!(env::var_os(&key), Some(OsString::from("value")));

    let result = panic::catch_unwind(panic::AssertUnwindSafe(|| {
        state.exit_scope(index + 1);
    }));
    assert!(result.is_ok(), "invalid scope exit should not panic");

    assert_eq!(env::var_os(&key), original);
    assert_eq!(state.depth(), 0);
    assert!(state.is_stack_empty());
    assert!(!state.has_lock());
    assert_env_lock_released();
}

/// Verifies corrupted scope state restores the environment and releases locks.
#[rstest]
#[case::corrupt_exit(CorruptionCase {
    test_name: "CORRUPT_EXIT",
    setup_guards: setup_single_guard,
    corrupt_state: apply_invalid_scope_exit,
    drop_guards: drop_guards_in_order,
    drop_message: "dropping guard after corruption should not panic",
})]
#[case::invalid_index_nested(CorruptionCase {
    test_name: "INVALID_INDEX_NESTED",
    setup_guards: setup_nested_guards,
    corrupt_state: apply_invalid_scope_exit,
    drop_guards: drop_guards_in_order,
    drop_message: "dropping guards after invalid scope exit should not panic",
})]
#[case::out_of_order_drop(CorruptionCase {
    test_name: "OUT_OF_ORDER_DROP",
    setup_guards: setup_nested_guards,
    corrupt_state: no_corruption,
    drop_guards: drop_guards_out_of_order,
    drop_message: "dropping outer guard out of order should not panic",
})]
#[serial]
fn scoped_env_recovers_from_corrupt_exit(#[case] case: CorruptionCase) {
    run_scoped_env_corruption_test(case.test_name, |key| {
        let original = env::var_os(key);
        let guards = (case.setup_guards)(key);
        let restored = (case.corrupt_state)();

        if restored {
            assert_eq!(env::var_os(key), original);
            assert_thread_state_reset();
        }

        let drop_result = panic::catch_unwind(panic::AssertUnwindSafe(|| {
            (case.drop_guards)(guards);
        }));
        assert!(drop_result.is_ok(), "{}", case.drop_message);
        assert_eq!(env::var_os(key), original);
        assert_thread_state_reset();
        assert_env_lock_released();
    });
}

/// Set an environment variable while holding `ENV_LOCK`.
fn set_env_var_locked(key: &OsStr, value: &OsStr) {
    let _guard = ENV_LOCK
        .lock()
        .unwrap_or_else(std::sync::PoisonError::into_inner);
    set_env_var_unlocked(key, value);
}

/// Set an environment variable when the caller already holds `ENV_LOCK`.
fn set_env_var_unlocked(key: &OsStr, value: &OsStr) {
    unsafe {
        // SAFETY: These helpers are only called from #[serial] tests, all
        // process environment mutations in this module are serialized via
        // ENV_LOCK, and callers such as RestoreEnv restore original values on
        // drop. Given these invariants and the caller-held lock, calling
        // env::set_var here is safe.
        env::set_var(key, value);
    }
}

/// Remove an environment variable when the caller already holds `ENV_LOCK`.
fn remove_env_var_unlocked(key: &OsStr) {
    unsafe {
        // SAFETY: These helpers are only called from #[serial] tests, all
        // process environment mutations in this module are serialized via
        // ENV_LOCK, and callers such as RestoreEnv restore original values on
        // drop. Given these invariants and the caller-held lock, calling
        // env::remove_var here is safe.
        env::remove_var(key);
    }
}

/// Assert that thread-local scoped environment state has fully reset.
fn assert_thread_state_reset() {
    THREAD_STATE.with(|cell| {
        let state = cell.borrow();
        assert_eq!(state.depth(), 0);
        assert!(state.is_stack_empty());
        assert!(!state.has_lock());
    });
}

/// Assert that `ENV_LOCK` can be acquired within a bounded timeout.
fn assert_env_lock_released() {
    let deadline = Instant::now() + Duration::from_secs(1);
    loop {
        match ENV_LOCK.try_lock() {
            Ok(guard) => {
                drop(guard);
                return;
            }
            Err(TryLockError::Poisoned(guard)) => {
                drop(guard);
                ENV_LOCK.clear_poison();
                return;
            }
            Err(TryLockError::WouldBlock) => {}
        }
        assert!(Instant::now() < deadline, "ENV_LOCK should be released");
        thread::yield_now();
    }
}

/// Verifies scoped environment application and restoration logs are emitted.
#[cfg(feature = "cluster-unit-tests")]
#[test]
#[serial]
fn logs_application_and_restoration() {
    let (logs, ()) = capture_info_logs(|| {
        let guard = ScopedEnv::apply(&[
            (String::from("OBS_ENV_APPLY"), Some(String::from("one"))),
            (String::from("OBS_ENV_CLEAR"), None),
        ]);
        drop(guard);
    });

    assert!(
        logs.iter()
            .any(|line| line.contains("applied scoped environment variables")),
        "expected application log entry, got {logs:?}"
    );
    assert!(
        logs.iter().any(|line| line.contains("OBS_ENV_APPLY=set")),
        "expected set entry, got {logs:?}"
    );
    assert!(
        logs.iter().any(|line| line.contains("OBS_ENV_CLEAR=unset")),
        "expected unset entry, got {logs:?}"
    );
    assert!(
        logs.iter()
            .any(|line| line.contains("restoring scoped environment variables")),
        "expected restoration log, got {logs:?}"
    );
}