From f019652d3eef1569dae55bac5cf8d93e4491fe90 Mon Sep 17 00:00:00 2001
From: kikijiki <dev@kikijiki.com>
Date: Sun, 29 Dec 2024 17:20:48 +0900
Subject: [PATCH] Add global setup/teardown hooks for test suites Fixes #251

---
 Cargo.toml              |  4 +++
 README.md               | 26 +++++++++++++++++
 src/macros.rs           | 28 ++++++++++++++++--
 tests/setup_teardown.rs | 65 +++++++++++++++++++++++++++++++++++++++++
 4 files changed, 121 insertions(+), 2 deletions(-)
 create mode 100644 tests/setup_teardown.rs

diff --git a/Cargo.toml b/Cargo.toml
index 802015934..526835d67 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -50,5 +50,9 @@ harness = false
 name = "integration"
 harness = true
 
+[[test]]
+name = "setup_teardown"
+harness = false
+
 [features]
 include-dir = ["dep:include_dir"]
diff --git a/README.md b/README.md
index 5eae3136e..174fa1471 100644
--- a/README.md
+++ b/README.md
@@ -189,6 +189,32 @@ Because the files don’t exist on disk, the test functions must accept their
 contents as either a `String` or a `Vec<u8>`. If the argument is not
 provided, the harness will panic at runtime.
 
+### Optional Setup and Teardown
+
+The harness macro supports optional setup and teardown functions that run before and after all tests:
+
+```rust
+fn setup() -> datatest_stable::Result<()> {
+    // Setup code here - runs before all tests.
+    Ok(())
+}
+
+fn teardown(exit_code: std::process::ExitCode) -> datatest_stable::Result<()> {
+    // Teardown code here - runs after all tests.
+    Ok(())
+}
+
+datatest_stable::harness! {
+    setup = setup,
+    teardown = teardown,
+    { test = my_test, root = "path/to/fixtures" },
+}
+```
+
+The setup function can be used to prepare test environment.
+The teardown function receives the exit code from test execution and can be used for cleanup.
+Both functions are optional and can be used independently.
+
 ### Conditionally embedding directories
 
 It is also possible to conditionally include directories at compile time via
diff --git a/src/macros.rs b/src/macros.rs
index 45d2907f7..fae56d15d 100644
--- a/src/macros.rs
+++ b/src/macros.rs
@@ -7,7 +7,11 @@
 /// a target](https://doc.rust-lang.org/cargo/reference/manifest.html#configuring-a-target)).
 #[macro_export]
 macro_rules! harness {
-    ( $( { $($args:tt)* } ),+ $(,)* ) => {
+    (
+        $(setup = $setup:path,)?
+        $(teardown = $teardown:path,)?
+        $( { $($args:tt)* } ),+ $(,)*
+    ) => {
         fn main() -> ::std::process::ExitCode {
             let mut requirements = Vec::new();
             use $crate::data_source_kinds::*;
@@ -17,7 +21,27 @@ macro_rules! harness {
                 $crate::harness_collect!(@gather_test requirements, { $($args)*, } => { });
             )+
 
-            $crate::runner(&requirements)
+            if !requirements.is_empty() {
+                $(
+                    if let Err(e) = $setup() {
+                        eprintln!("Setup failed: {}", e);
+                        return ::std::process::ExitCode::FAILURE;
+                    }
+                )?
+
+                let exit_code = $crate::runner(&requirements);
+
+                $(
+                    if let Err(e) = $teardown(exit_code) {
+                        eprintln!("Teardown failed: {}", e);
+                        return ::std::process::ExitCode::FAILURE;
+                    }
+                )?
+
+                exit_code
+            } else {
+                ::std::process::ExitCode::SUCCESS
+            }
         }
     };
     ( $( $name:path, $root:expr, $pattern:expr ),+ $(,)* ) => {
diff --git a/tests/setup_teardown.rs b/tests/setup_teardown.rs
new file mode 100644
index 000000000..31996d91d
--- /dev/null
+++ b/tests/setup_teardown.rs
@@ -0,0 +1,65 @@
+// Copyright (c) The datatest-stable Contributors
+// SPDX-License-Identifier: MIT OR Apache-2.0
+
+use camino::Utf8Path;
+use datatest_stable::Result;
+use std::sync::atomic::{AtomicUsize, Ordering};
+use std::path::Path;
+
+static SETUP_COUNTER: AtomicUsize = AtomicUsize::new(0);
+static TEARDOWN_COUNTER: AtomicUsize = AtomicUsize::new(0);
+static TEST_ARTIFACT_COUNTER: AtomicUsize = AtomicUsize::new(0);
+static TEST_ARTIFACT_UTF8_COUNTER: AtomicUsize = AtomicUsize::new(0);
+
+fn test_setup() -> Result<()> {
+    SETUP_COUNTER.fetch_add(1, Ordering::SeqCst);
+    println!("Setup called");
+    Ok(())
+}
+
+fn test_teardown(exit_code: std::process::ExitCode) -> Result<()> {
+    TEARDOWN_COUNTER.fetch_add(1, Ordering::SeqCst);
+    if exit_code == std::process::ExitCode::SUCCESS {
+        println!("Teardown called after success");
+    } else {
+        println!("Teardown called after failure");
+    }
+    Ok(())
+}
+
+fn test_artifact(_path: &Path) -> Result<()> {
+    TEST_ARTIFACT_COUNTER.fetch_add(1, Ordering::SeqCst);
+    Ok(())
+}
+
+fn test_artifact_utf8(_path: &Utf8Path) -> Result<()> {
+    TEST_ARTIFACT_UTF8_COUNTER.fetch_add(1, Ordering::SeqCst);
+    Ok(())
+}
+
+datatest_stable::harness! {
+    setup = test_setup,
+    teardown = test_teardown,
+    {
+        test = test_artifact,
+        root = "tests/files",
+    },
+    {
+        test = test_artifact_utf8,
+        root = "tests/files",
+    }
+}
+
+#[test]
+fn verify_counters() {
+    // Setup should be called exactly once.
+    assert_eq!(SETUP_COUNTER.load(Ordering::SeqCst), 1);
+
+    // Teardown should be called exactly once.
+    assert_eq!(TEARDOWN_COUNTER.load(Ordering::SeqCst), 1);
+
+    // Each test function should be called once per file in tests/files.
+    let file_count = fs::read_dir("tests/files").unwrap().count();
+    assert_eq!(TEST_ARTIFACT_COUNTER.load(Ordering::SeqCst), file_count);
+    assert_eq!(TEST_ARTIFACT_UTF8_COUNTER.load(Ordering::SeqCst), file_count);
+}