Fix support for import autoinstrument and TestMain scopeagent.Run

init.go

@@ -6,6 +6,7 @@ import (
 	"os"
 	"os/signal"
 	"runtime"
+	"sync"
 	"syscall"
 	"testing"
 
@@ -15,10 +16,19 @@ import (
 	scopetesting "go.undefinedlabs.com/scopeagent/instrumentation/testing"
 )
 
-var defaultAgent *agent.Agent
+var (
+	defaultAgent *agent.Agent
+	runningMutex sync.RWMutex
+	running      bool
+)
 
 // Helper function to run a `testing.M` object and gracefully stopping the agent afterwards
 func Run(m *testing.M, opts ...agent.Option) int {
+	if getRunningFlag() {
+		return m.Run()
+	}
+	setRunningFlag(true)
+	defer setRunningFlag(false)
 	opts = append(opts, agent.WithTestingModeEnabled())
 	newAgent, err := agent.NewAgent(opts...)
 	if err != nil {
@@ -130,3 +140,14 @@ func StartBenchmark(b *testing.B, benchFunc func(b *testing.B)) {
 	pc, _, _, _ := runtime.Caller(1)
 	scopetesting.StartBenchmark(b, pc, benchFunc)
 }
+
+func setRunningFlag(value bool) {
+	runningMutex.Lock()
+	defer runningMutex.Unlock()
+	running = value
+}
+func getRunningFlag() bool {
+	runningMutex.RLock()
+	defer runningMutex.RUnlock()
+	return running
+}