From 985e770c27a9b9a53400b827649e482d22781e10 Mon Sep 17 00:00:00 2001
From: ppacher <ppacher@users.noreply.github.com>
Date: Mon, 6 Apr 2020 12:35:51 +0200
Subject: [PATCH 1/2] Add support for start and stop hooks during test setup

---
 core/pmtesting/testing.go | 21 +++++++++++++++++++++
 1 file changed, 21 insertions(+)

diff --git a/core/pmtesting/testing.go b/core/pmtesting/testing.go
index 04779939..49445d52 100644
--- a/core/pmtesting/testing.go
+++ b/core/pmtesting/testing.go
@@ -41,8 +41,17 @@ func init() {
 	flag.BoolVar(&printStackOnExit, "print-stack-on-exit", false, "prints the stack before of shutting down")
 }
 
+type TestHookFunc func() error
+
 // TestMain provides a simple unit test setup routine.
 func TestMain(m *testing.M, module *modules.Module) {
+	TestMainWithHooks(m, module, nil, nil)
+}
+
+// TestMainWithHooks provides a simple unit test setup routine and calls
+// afterStartFn after modules have started and beforeStopFn before modules
+// are shutdown.
+func TestMainWithHooks(m *testing.M, module *modules.Module, afterStartFn, beforeStopFn TestHookFunc) {
 	// enable module for testing
 	module.Enable()
 
@@ -72,10 +81,22 @@ func TestMain(m *testing.M, module *modules.Module) {
 		fmt.Fprintf(os.Stderr, "failed to setup test: %s\n", err)
 		exitCode = 1
 	} else {
+		if afterStartFn != nil {
+			if err := afterStartFn(); err != nil {
+				fmt.Fprintf(os.Stderr, "failed to run test start hook: %s\n", err)
+				exitCode = 1
+			}
+		}
 		// run tests
 		exitCode = m.Run()
 	}
 
+	if beforeStopFn != nil {
+		if err := beforeStopFn(); err != nil {
+			fmt.Fprintf(os.Stderr, "failed to run test shutdown hook: %s\n", err)
+		}
+	}
+
 	// shutdown
 	_ = modules.Shutdown()
 	if modules.GetExitStatusCode() != 0 {

From 81d2b6309dd4e077aace19e6c2ab4955cdb5909c Mon Sep 17 00:00:00 2001
From: ppacher <ppacher@users.noreply.github.com>
Date: Mon, 6 Apr 2020 13:11:55 +0200
Subject: [PATCH 2/2] Only run tests if start hook succeeds

---
 core/pmtesting/testing.go | 9 +++++++--
 1 file changed, 7 insertions(+), 2 deletions(-)

diff --git a/core/pmtesting/testing.go b/core/pmtesting/testing.go
index 49445d52..79925811 100644
--- a/core/pmtesting/testing.go
+++ b/core/pmtesting/testing.go
@@ -81,14 +81,19 @@ func TestMainWithHooks(m *testing.M, module *modules.Module, afterStartFn, befor
 		fmt.Fprintf(os.Stderr, "failed to setup test: %s\n", err)
 		exitCode = 1
 	} else {
+		runTests := true
 		if afterStartFn != nil {
 			if err := afterStartFn(); err != nil {
 				fmt.Fprintf(os.Stderr, "failed to run test start hook: %s\n", err)
+				runTests = false
 				exitCode = 1
 			}
 		}
-		// run tests
-		exitCode = m.Run()
+
+		if runTests {
+			// run tests
+			exitCode = m.Run()
+		}
 	}
 
 	if beforeStopFn != nil {