pdk-test 1.7.0

// Copyright (c) 2026, Salesforce, Inc.,
// All rights reserved.
// For full license text, see the LICENSE.txt file

use std::cell::Cell;
use std::future::Future;
use std::path::{Path, PathBuf};
use std::rc::Rc;
use std::sync::{Mutex, Once, OnceLock, PoisonError};

use crate::TestError;

// Initialize logging once for test runs.
//
// Logs pdk_test=info by default or uses RUST_LOG value if set.
fn init_logger() {
    static INIT: Once = Once::new();
    INIT.call_once(|| {
        let env = env_logger::Env::default().default_filter_or("pdk_test=info");
        let _ = env_logger::Builder::from_env(env).is_test(true).try_init();
    });
}

fn init_target_dir() -> PathBuf {
    let manifest_dir =
        std::env::var("CARGO_MANIFEST_DIR").expect("CARGO_MANIFEST_DIR env variable not found");

    Path::new(&manifest_dir)
        .ancestors()
        .map(|parent| parent.join("target"))
        .find(|candidate| candidate.is_dir())
        .expect("Cargo 'target/' directory not found.")
}

fn target_dir() -> &'static Path {
    static TARGET_DIR: OnceLock<PathBuf> = OnceLock::new();

    TARGET_DIR.get_or_init(init_target_dir)
}

/// Trait to check if a test result is successful.
pub trait TestResult {
    fn is_successful(&self) -> bool;
}

impl TestResult for () {
    fn is_successful(&self) -> bool {
        true
    }
}

impl<T, E> TestResult for Result<T, E> {
    fn is_successful(&self) -> bool {
        self.is_ok()
    }
}

tokio::task_local! {
    static CURRENT_TEST: Rc<Test>;
}

/// Struct to handle test lifecycle.
/// Provides access to test information and lifecycle methods.
#[derive(Debug)]
pub struct Test {
    name: String,
    module: String,
    target_dir: PathBuf,
    success: Cell<bool>,
    force_upload: bool,
}

impl Test {
    /// Get the current test instance.
    pub(crate) fn current() -> Result<Rc<Self>, TestError> {
        CURRENT_TEST.try_with(Rc::clone).map_err(|_| {
            TestError::Startup(
                "Test not started. \
                    Test functions must be decorated with #[pdk_test] attribute."
                    .to_string(),
            )
        })
    }

    /// Check if the test was successful.
    pub(crate) fn is_success(&self) -> bool {
        self.success.get()
    }

    /// Get the test builder.
    pub fn builder() -> TestBuilder {
        TestBuilder::new()
    }

    /// Get the test name.
    pub fn name(&self) -> &str {
        &self.name
    }

    /// Get the test module.
    pub fn module(&self) -> &str {
        &self.module
    }

    /// Check if the test should force upload.
    pub fn force_upload(&self) -> bool {
        self.force_upload
    }

    /// Get the test target directory.
    pub(crate) fn target_dir(&self) -> &Path {
        &self.target_dir
    }

    fn try_remove_target_dir(&self) -> std::io::Result<()> {
        // Remove the test directory
        std::fs::remove_dir_all(self.target_dir())?;

        let parent = self.target_dir().parent().unwrap();

        // If the module directory is empty, remove it too
        if std::fs::read_dir(parent)?.next().is_none() {
            std::fs::remove_dir(parent)?;
        }

        Ok(())
    }

    fn remove_target_dir(&self) {
        if let Err(e) = self.try_remove_target_dir() {
            log::error!(
                "Unable to remove test directory {}: {e}",
                self.target_dir().display()
            );
        }
    }

    /// Run the test.
    pub fn run<F>(self, body: F) -> F::Output
    where
        F: Future,
        F::Output: TestResult,
    {
        init_logger();
        static LOCK: Mutex<()> = Mutex::new(());

        let _guard = LOCK.lock().unwrap_or_else(PoisonError::into_inner);

        std::fs::create_dir_all(self.target_dir()).expect("Unable to create test directory");

        // Scope current test as task local
        let body = CURRENT_TEST.scope(Rc::new(self), async move {
            let result = body.await;

            if result.is_successful() {
                // Perform post-success actions
                CURRENT_TEST.with(|test| {
                    test.remove_target_dir();
                    test.success.set(true);
                });
            }

            result
        });

        tokio::runtime::Builder::new_multi_thread()
            .worker_threads(2usize)
            .enable_all()
            .build()
            .expect("Failed building the Runtime")
            .block_on(body)
    }
}

/// Struct to build a test.
#[derive(Default)]
pub struct TestBuilder {
    name: Option<String>,
    module: Option<String>,

    #[allow(unused)]
    force_upload: bool,
}

impl TestBuilder {
    fn new() -> Self {
        Self {
            name: None,
            module: None,
            force_upload: false,
        }
    }

    /// Set the test module.
    pub fn module<T: Into<String>>(self, module: T) -> Self {
        Self {
            module: Some(module.into()),
            ..self
        }
    }

    /// Set the test name.
    pub fn name<T: Into<String>>(self, name: T) -> Self {
        Self {
            name: Some(name.into()),
            ..self
        }
    }

    /// Set the test force upload.
    pub fn force_upload(self) -> Self {
        Self {
            force_upload: true,
            ..self
        }
    }

    /// Build test instance.
    pub fn build(self) -> Test {
        let name = self.name.expect("Test name is required.");
        let module = self.module.expect("Test module is required");
        let target_dir = target_dir().join("pdk-test").join(&module).join(&name);

        Test {
            name,
            module,
            target_dir,
            force_upload: self.force_upload,
            success: Cell::new(false),
        }
    }
}

#[cfg(test)]
mod tests {
    use crate::handle::handle;
    use std::path::PathBuf;

    use super::Test;

    struct TestProvider {
        test: Option<Test>,
        target: Option<PathBuf>,
    }

    impl TestProvider {
        fn new() -> Self {
            Self {
                test: Some(
                    Test::builder()
                        .module(handle("bar"))
                        .name(handle("foo"))
                        .build(),
                ),
                target: None,
            }
        }

        fn take(&mut self) -> Test {
            let test = self.test.take().expect("Already taken test.");
            self.target = Some(test.target_dir().into());
            test
        }
    }

    // Drop ensures dir cleanup
    impl Drop for TestProvider {
        fn drop(&mut self) {
            let target = self
                .target
                .as_deref()
                .or_else(|| self.test.as_ref().map(Test::target_dir))
                .unwrap();

            let parent = target.parent().unwrap();
            if parent.exists() {
                std::fs::remove_dir_all(parent).unwrap();
            }
        }
    }

    #[test]
    fn dir_path() {
        let mut provider = TestProvider::new();
        let test = provider.take();

        let module = test.module();
        let name = test.name();

        let target_dir = test.target_dir().to_str().unwrap();

        assert!(target_dir.ends_with(&format!("/target/pdk-test/{module}/{name}")));
        assert!(target_dir.starts_with(super::target_dir().to_str().unwrap()))
    }

    #[test]
    fn retain_dir_on_fail() {
        let mut provider = TestProvider::new();
        let test = provider.take();

        let target_dir = test.target_dir().to_owned();

        let _ = test.run(async { Err::<(), ()>(()) });

        let target_dir_exists = target_dir.exists();

        assert!(target_dir_exists);
    }

    #[test]
    fn remove_dir_on_success() {
        let mut provider = TestProvider::new();
        let test = provider.take();

        let target_dir = test.target_dir().to_owned();

        test.run(async {
            assert!(target_dir.exists());
        });

        assert!(
            !target_dir.exists(),
            "Target dir {} not deleted",
            target_dir.display()
        );
    }

    #[test]
    fn test_not_created() {
        assert!(Test::current().is_err());
    }
}