From a01bd0fa7efa89ef0c5aab2047e4f896ead4ecf7 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Maximilian=20Pa=C3=9F?=
 <22845248+mpass99@users.noreply.github.com>
Date: Mon, 4 Sep 2023 11:20:32 +0200
Subject: [PATCH] Provide Memory Leak Test Suite

by adding an assertion about the number of Goroutines to the unit tests.
---
 tests/constants.go |  1 +
 tests/util.go      | 50 +++++++++++++++++++++++++++++++++++++++++++++-
 2 files changed, 50 insertions(+), 1 deletion(-)

diff --git a/tests/constants.go b/tests/constants.go
index 6d41f2a..457a151 100644
--- a/tests/constants.go
+++ b/tests/constants.go
@@ -25,6 +25,7 @@ const (
 	AnotherRunnerID               = AnotherEnvironmentIDAsString + "-" + AnotherUUID
 	DefaultExecutionID            = "s0m3-3x3cu710n-1d"
 	DefaultMockID                 = "m0ck-1d"
+	TinyTimeout                   = 10 * time.Millisecond
 	ShortTimeout                  = 100 * time.Millisecond
 	DefaultTestTimeout            = 10 * time.Minute
 )
diff --git a/tests/util.go b/tests/util.go
index f72b4e4..6b4ec20 100644
--- a/tests/util.go
+++ b/tests/util.go
@@ -1,6 +1,15 @@
 package tests
 
-import "time"
+import (
+	"bytes"
+	"context"
+	"github.com/stretchr/testify/suite"
+	"io"
+	"os"
+	"runtime"
+	"runtime/pprof"
+	"time"
+)
 
 // ChannelReceivesSomething waits timeout seconds for something to be received from channel ch.
 // If something is received, it returns true. If the timeout expires without receiving anything, it return false.
@@ -12,3 +21,42 @@ func ChannelReceivesSomething(ch chan bool, timeout time.Duration) bool {
 		return false
 	}
 }
+
+// MemoryLeakTestSuite adds an assertion for checking Goroutine leaks.
+// Be aware not to overwrite the SetupTest or TearDownTest function!
+type MemoryLeakTestSuite struct {
+	suite.Suite
+	ExpectedGoroutingIncrease int
+	TestCtx                   context.Context
+	testCtxCancel             context.CancelFunc
+	goroutineCountBefore      int
+	goroutinesBefore          *bytes.Buffer
+}
+
+func (s *MemoryLeakTestSuite) SetupTest() {
+	s.goroutineCountBefore = runtime.NumGoroutine()
+	s.ExpectedGoroutingIncrease = 0
+	s.goroutinesBefore = &bytes.Buffer{}
+
+	err := pprof.Lookup("goroutine").WriteTo(s.goroutinesBefore, 1)
+	s.NoError(err)
+
+	ctx, cancel := context.WithCancel(context.Background())
+	s.TestCtx = ctx
+	s.testCtxCancel = cancel
+}
+
+func (s *MemoryLeakTestSuite) TearDownTest() {
+	s.testCtxCancel()
+	runtime.Gosched()         // Flush done Goroutines
+	<-time.After(TinyTimeout) // Just to make sure
+	goroutinesAfter := runtime.NumGoroutine()
+	s.Equal(s.goroutineCountBefore+s.ExpectedGoroutingIncrease, goroutinesAfter)
+
+	if s.goroutineCountBefore+s.ExpectedGoroutingIncrease != goroutinesAfter {
+		_, err := io.Copy(os.Stderr, s.goroutinesBefore)
+		s.NoError(err)
+		err = pprof.Lookup("goroutine").WriteTo(os.Stderr, 1)
+		s.NoError(err)
+	}
+}