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] 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 {