[NOD-299] Add waitgroup to wait for all spawns to complete before calling teardown (#385)

* [NOD-299] Add waitgroup to wait for all `spawn`s to complete before calling teardown

* [NOD-299] Restore spawn on teardown + mark spawn done in the correct thread

blockdag/test_utils.go

@@ -6,6 +6,7 @@ import (
 	"fmt"
 	"os"
 	"path/filepath"
+	"sync"
 
 	"github.com/daglabs/btcd/util/subnetworkid"
 
@@ -60,6 +61,18 @@ func DAGSetup(dbName string, config Config) (*BlockDAG, func(), error) {
 
 	var teardown func()
 
+	// To make sure that the teardown function is not called before any goroutines finished to run -
+	// overwrite `spawn` to count the number of running goroutines
+	spawnWaitGroup := sync.WaitGroup{}
+	realSpawn := spawn
+	spawn = func(f func()) {
+		spawnWaitGroup.Add(1)
+		realSpawn(func() {
+			f()
+			spawnWaitGroup.Done()
+		})
+	}
+
 	if config.DB == nil {
 		// Create the root directory for test databases.
 		if !fileExists(testDbRoot) {
@@ -81,6 +94,8 @@ func DAGSetup(dbName string, config Config) (*BlockDAG, func(), error) {
 		// Setup a teardown function for cleaning up.  This function is
 		// returned to the caller to be invoked when it is done testing.
 		teardown = func() {
+			spawnWaitGroup.Wait()
+			spawn = realSpawn
 			config.DB.Close()
 			os.RemoveAll(dbPath)
 			os.RemoveAll(testDbRoot)