diff --git a/init.go b/init.go index 930d4abc..fb3ae83d 100644 --- a/init.go +++ b/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 +}