From 83680934e5f04919551ca99645054c470b4bc062 Mon Sep 17 00:00:00 2001
From: Peter Elliott <pelliott@ualberta.ca>
Date: Sun, 29 Aug 2021 11:55:05 -0700
Subject: [PATCH] LibTest: Add TEST_SETUP macro that runs before all test cases

---
 Userland/Libraries/LibTest/TestCase.h    | 10 +++++++++-
 Userland/Libraries/LibTest/TestSuite.cpp | 10 ++++++++++
 Userland/Libraries/LibTest/TestSuite.h   |  3 +++
 3 files changed, 22 insertions(+), 1 deletion(-)

diff --git a/Userland/Libraries/LibTest/TestCase.h b/Userland/Libraries/LibTest/TestCase.h
index d0fb713567f..866bba5d2ec 100644
--- a/Userland/Libraries/LibTest/TestCase.h
+++ b/Userland/Libraries/LibTest/TestCase.h
@@ -39,9 +39,17 @@ private:
 
 // Helper to hide implementation of TestSuite from users
 void add_test_case_to_suite(const NonnullRefPtr<TestCase>& test_case);
-
+void set_suite_setup_function(Function<void()> setup);
 }
 
+#define TEST_SETUP                                                  \
+    static void __setup();                                          \
+    struct __setup_type {                                           \
+        __setup_type() { Test::set_suite_setup_function(__setup); } \
+    };                                                              \
+    static struct __setup_type __setup_type;                        \
+    static void __setup()
+
 #define __TESTCASE_FUNC(x) __test_##x
 #define __TESTCASE_TYPE(x) __TestCase_##x
 
diff --git a/Userland/Libraries/LibTest/TestSuite.cpp b/Userland/Libraries/LibTest/TestSuite.cpp
index ac2b93b7de1..e6923040ea1 100644
--- a/Userland/Libraries/LibTest/TestSuite.cpp
+++ b/Userland/Libraries/LibTest/TestSuite.cpp
@@ -7,6 +7,7 @@
 
 #include <LibTest/Macros.h> // intentionally first -- we redefine VERIFY and friends in here
 
+#include <AK/Function.h>
 #include <LibCore/ArgsParser.h>
 #include <LibTest/TestSuite.h>
 #include <stdlib.h>
@@ -49,6 +50,12 @@ void add_test_case_to_suite(const NonnullRefPtr<TestCase>& test_case)
     TestSuite::the().add_case(test_case);
 }
 
+// Declared in TestCase.h
+void set_suite_setup_function(Function<void()> setup)
+{
+    TestSuite::the().set_suite_setup(move(setup));
+}
+
 int TestSuite::main(const String& suite_name, int argc, char** argv)
 {
     m_suite_name = suite_name;
@@ -66,6 +73,9 @@ int TestSuite::main(const String& suite_name, int argc, char** argv)
     args_parser.add_positional_argument(search_string, "Only run matching cases.", "pattern", Core::ArgsParser::Required::No);
     args_parser.parse(argc, argv);
 
+    if (m_setup)
+        m_setup();
+
     const auto& matching_tests = find_cases(search_string, !do_benchmarks_only, !do_tests_only);
 
     if (do_list_cases) {
diff --git a/Userland/Libraries/LibTest/TestSuite.h b/Userland/Libraries/LibTest/TestSuite.h
index 694b6d1a696..4cc11425140 100644
--- a/Userland/Libraries/LibTest/TestSuite.h
+++ b/Userland/Libraries/LibTest/TestSuite.h
@@ -43,6 +43,8 @@ public:
 
     void current_test_case_did_fail() { m_current_test_case_passed = false; }
 
+    void set_suite_setup(Function<void()> setup) { m_setup = move(setup); }
+
 private:
     static TestSuite* s_global;
     NonnullRefPtrVector<TestCase> m_cases;
@@ -50,6 +52,7 @@ private:
     u64 m_benchtime = 0;
     String m_suite_name;
     bool m_current_test_case_passed = true;
+    Function<void()> m_setup;
 };
 
 }