Add testhelper

cmd/check_example/main.go

@@ -1,42 +1,34 @@
 package main
 
 import (
-	"math/rand"
-	"os"
-	"time"
-
 	"github.com/NETWAYS/go-check"
 	log "github.com/sirupsen/logrus"
+	"os"
 )
 
-func initRand() int {
-	rand.Seed(time.Now().UnixNano())
-	min := 1
-	max := 100
-	return rand.Intn(max-min+1) + min
-}
-
 func main() {
+	defer check.CatchPanic()
 	flags := check.NewFlags()
 	flags.Name = "check_test"
 	flags.Readme = `Test Plugin`
 	flags.Version = "1.0.0"
 
+	value := flags.Set.IntP("value", "t", 10, "test value")
 	warning := flags.Set.IntP("warning", "w", 20, "warning threshold")
 	critical := flags.Set.IntP("critical", "c", 50, "critical threshold")
 
-	ranNum := initRand()
+	// value should be calculated
 
 	flags.Parse(os.Args[1:])
 	flags.SetupLogging()
 
 	log.Info("Start logging")
 
-	if *critical < ranNum {
-		check.Exit(check.Critical, "value is %d", ranNum)
-	} else if *warning < ranNum {
-		check.Exit(check.Warning, "value is %d", ranNum)
+	if *value > *critical {
+		check.Exit(check.Critical, "value is %d", *value)
+	} else if *value > *warning {
+		check.Exit(check.Warning, "value is %d", *value)
 	} else {
-		check.Exit(check.OK, "value is %d", ranNum)
+		check.Exit(check.OK, "value is %d", *value)
 	}
 }

cmd/check_example/main_test.go

@@ -0,0 +1,25 @@
+package main
+
+import (
+	"github.com/NETWAYS/go-check/testhelper"
+	"github.com/stretchr/testify/assert"
+	"os"
+	"regexp"
+	"testing"
+)
+
+func TestMyMain(t *testing.T) {
+	stdout := testhelper.RunMainTest(main, "--help")
+	assert.Regexp(t, regexp.MustCompile(`would exit with code 3`), stdout)
+
+	stdout = testhelper.RunMainTest(main, "--warning", "20")
+	assert.Regexp(t, regexp.MustCompile(`^OK - value is 10\nwould exit with code 0\n$`), stdout)
+
+	stdout = testhelper.RunMainTest(main, "--warning", "10", "--value", "11")
+	assert.Regexp(t, regexp.MustCompile(`^WARNING - value is 11\nwould exit with code 1\n$`), stdout)
+}
+
+func TestMain(m *testing.M) {
+	testhelper.EnableTestMode()
+	os.Exit(m.Run())
+}

testhelper/for_main.go

@@ -0,0 +1,58 @@
+package testhelper
+
+import (
+	"bytes"
+	"github.com/NETWAYS/go-check"
+	"io"
+	"os"
+)
+
+// Execute main function from a main package, while capturing its stdout
+//
+// You will need to pass `main`, since the function won't have access to the package's namespace
+func RunMainTest(f func(), args ...string) string {
+	base := []string{"check_with_go_test"}
+	origArgs := os.Args
+	os.Args = append(base, args...)
+	stdout := CaptureStdout(f)
+	os.Args = origArgs
+	return stdout
+}
+
+// Enable test mode by disabling the default exit behavior, so go test won't fail with plugin states
+func EnableTestMode() {
+	// disable actual exit
+	check.AllowExit = false
+
+	// disable stack trace for the example
+	check.PrintStack = false
+}
+
+// Disable test mode behavior again
+//
+// Optional after testing has been done
+func DisableTestMode() {
+	// disable actual exit
+	check.AllowExit = true
+
+	// disable stack trace for the example
+	check.PrintStack = true
+}
+
+// Capture the output of the go program while running function f
+//
+// Source https://gist.github.com/mindscratch/0faa78bd3c0005d080bf
+func CaptureStdout(f func()) string {
+	old := os.Stdout
+	r, w, _ := os.Pipe()
+	os.Stdout = w
+
+	f()
+
+	w.Close()
+	os.Stdout = old
+
+	var buf bytes.Buffer
+	_, _ = io.Copy(&buf, r)
+	return buf.String()
+}