watchtower/wtdb: use T.TempDir to create temporary test directory

Signed-off-by: Eng Zer Jun <engzerjun@gmail.com>

watchtower/wtdb/client_db_test.go

@@ -4,9 +4,7 @@ import (
 	"bytes"
 	crand "crypto/rand"
 	"io"
-	"io/ioutil"
 	"net"
-	"os"
 	"reflect"
 	"testing"
 
@@ -20,24 +18,23 @@ import (
 	"github.com/lightningnetwork/lnd/watchtower/wtpolicy"
 )
 
-// clientDBInit is a closure used to initialize a wtclient.DB instance its
+// clientDBInit is a closure used to initialize a wtclient.DB instance.
-// cleanup function.
+type clientDBInit func(t *testing.T) wtclient.DB
-type clientDBInit func(t *testing.T) (wtclient.DB, func())
 
 type clientDBHarness struct {
 	t  *testing.T
 	db wtclient.DB
 }
 
-func newClientDBHarness(t *testing.T, init clientDBInit) (*clientDBHarness, func()) {
+func newClientDBHarness(t *testing.T, init clientDBInit) *clientDBHarness {
-	db, cleanup := init(t)
+	db := init(t)
 
 	h := &clientDBHarness{
 		t:  t,
 		db: db,
 	}
 
-	return h, cleanup
+	return h
 }
 
 func (h *clientDBHarness) insertSession(session *wtdb.ClientSession, expErr error) {
@@ -778,55 +775,40 @@ func TestClientDB(t *testing.T) {
 	}{
 		{
 			name: "fresh clientdb",
-			init: func(t *testing.T) (wtclient.DB, func()) {
+			init: func(t *testing.T) wtclient.DB {
-				path, err := ioutil.TempDir("", "clientdb")
-				if err != nil {
-					t.Fatalf("unable to make temp dir: %v",
-						err)
-				}
-
 				bdb, err := wtdb.NewBoltBackendCreator(
-					true, path, "wtclient.db",
+					true, t.TempDir(), "wtclient.db",
 				)(dbCfg)
 				if err != nil {
-					os.RemoveAll(path)
 					t.Fatalf("unable to open db: %v", err)
 				}
 
 				db, err := wtdb.OpenClientDB(bdb)
 				if err != nil {
-					os.RemoveAll(path)
 					t.Fatalf("unable to open db: %v", err)
 				}
 
-				cleanup := func() {
+				t.Cleanup(func() {
 					db.Close()
-					os.RemoveAll(path)
+				})
-				}
 
-				return db, cleanup
+				return db
 			},
 		},
 		{
 			name: "reopened clientdb",
-			init: func(t *testing.T) (wtclient.DB, func()) {
+			init: func(t *testing.T) wtclient.DB {
-				path, err := ioutil.TempDir("", "clientdb")
+				path := t.TempDir()
-				if err != nil {
-					t.Fatalf("unable to make temp dir: %v",
-						err)
-				}
 
 				bdb, err := wtdb.NewBoltBackendCreator(
 					true, path, "wtclient.db",
 				)(dbCfg)
 				if err != nil {
-					os.RemoveAll(path)
 					t.Fatalf("unable to open db: %v", err)
 				}
 
 				db, err := wtdb.OpenClientDB(bdb)
 				if err != nil {
-					os.RemoveAll(path)
 					t.Fatalf("unable to open db: %v", err)
 				}
 				db.Close()
@@ -835,28 +817,25 @@ func TestClientDB(t *testing.T) {
 					true, path, "wtclient.db",
 				)(dbCfg)
 				if err != nil {
-					os.RemoveAll(path)
 					t.Fatalf("unable to open db: %v", err)
 				}
 
 				db, err = wtdb.OpenClientDB(bdb)
 				if err != nil {
-					os.RemoveAll(path)
 					t.Fatalf("unable to reopen db: %v", err)
 				}
 
-				cleanup := func() {
+				t.Cleanup(func() {
 					db.Close()
-					os.RemoveAll(path)
+				})
-				}
 
-				return db, cleanup
+				return db
 			},
 		},
 		{
 			name: "mock",
-			init: func(t *testing.T) (wtclient.DB, func()) {
+			init: func(t *testing.T) wtclient.DB {
-				return wtmock.NewClientDB(), func() {}
+				return wtmock.NewClientDB()
 			},
 		},
 	}
@@ -902,10 +881,7 @@ func TestClientDB(t *testing.T) {
 
 			for _, test := range tests {
 				t.Run(test.name, func(t *testing.T) {
-					h, cleanup := newClientDBHarness(
+					h := newClientDBHarness(t, db.init)
-						t, db.init,
-					)
-					defer cleanup()
 
 					test.run(h)
 				})

watchtower/wtdb/tower_db_test.go

@@ -3,8 +3,6 @@ package wtdb_test
 import (
 	"bytes"
 	"encoding/binary"
-	"io/ioutil"
-	"os"
 	"reflect"
 	"testing"
 
@@ -22,9 +20,8 @@ var (
 	testBlob = make([]byte, blob.Size(blob.TypeAltruistCommit))
 )
 
-// dbInit is a closure used to initialize a watchtower.DB instance and its
+// dbInit is a closure used to initialize a watchtower.DB instance.
-// cleanup function.
+type dbInit func(*testing.T) watchtower.DB
-type dbInit func(*testing.T) (watchtower.DB, func())
 
 // towerDBHarness holds the resources required to execute the tower db tests.
 type towerDBHarness struct {
@@ -34,15 +31,15 @@ type towerDBHarness struct {
 
 // newTowerDBHarness initializes a fresh test harness for testing watchtower.DB
 // implementations.
-func newTowerDBHarness(t *testing.T, init dbInit) (*towerDBHarness, func()) {
+func newTowerDBHarness(t *testing.T, init dbInit) *towerDBHarness {
-	db, cleanup := init(t)
+	db := init(t)
 
 	h := &towerDBHarness{
 		t:  t,
 		db: db,
 	}
 
-	return h, cleanup
+	return h
 }
 
 // insertSession attempts to isnert the passed session and asserts that the
@@ -637,55 +634,42 @@ func TestTowerDB(t *testing.T) {
 	}{
 		{
 			name: "fresh boltdb",
-			init: func(t *testing.T) (watchtower.DB, func()) {
+			init: func(t *testing.T) watchtower.DB {
-				path, err := ioutil.TempDir("", "towerdb")
+				path := t.TempDir()
-				if err != nil {
-					t.Fatalf("unable to make temp dir: %v",
-						err)
-				}
 
 				bdb, err := wtdb.NewBoltBackendCreator(
 					true, path, "watchtower.db",
 				)(dbCfg)
 				if err != nil {
-					os.RemoveAll(path)
 					t.Fatalf("unable to open db: %v", err)
 				}
 
 				db, err := wtdb.OpenTowerDB(bdb)
 				if err != nil {
-					os.RemoveAll(path)
 					t.Fatalf("unable to open db: %v", err)
 				}
 
-				cleanup := func() {
+				t.Cleanup(func() {
 					db.Close()
-					os.RemoveAll(path)
+				})
-				}
 
-				return db, cleanup
+				return db
 			},
 		},
 		{
 			name: "reopened boltdb",
-			init: func(t *testing.T) (watchtower.DB, func()) {
+			init: func(t *testing.T) watchtower.DB {
-				path, err := ioutil.TempDir("", "towerdb")
+				path := t.TempDir()
-				if err != nil {
-					t.Fatalf("unable to make temp dir: %v",
-						err)
-				}
 
 				bdb, err := wtdb.NewBoltBackendCreator(
 					true, path, "watchtower.db",
 				)(dbCfg)
 				if err != nil {
-					os.RemoveAll(path)
 					t.Fatalf("unable to open db: %v", err)
 				}
 
 				db, err := wtdb.OpenTowerDB(bdb)
 				if err != nil {
-					os.RemoveAll(path)
 					t.Fatalf("unable to open db: %v", err)
 				}
 				db.Close()
@@ -697,28 +681,25 @@ func TestTowerDB(t *testing.T) {
 					true, path, "watchtower.db",
 				)(dbCfg)
 				if err != nil {
-					os.RemoveAll(path)
 					t.Fatalf("unable to open db: %v", err)
 				}
 
 				db, err = wtdb.OpenTowerDB(bdb)
 				if err != nil {
-					os.RemoveAll(path)
 					t.Fatalf("unable to open db: %v", err)
 				}
 
-				cleanup := func() {
+				t.Cleanup(func() {
 					db.Close()
-					os.RemoveAll(path)
+				})
-				}
 
-				return db, cleanup
+				return db
 			},
 		},
 		{
 			name: "mock",
-			init: func(t *testing.T) (watchtower.DB, func()) {
+			init: func(t *testing.T) watchtower.DB {
-				return wtmock.NewTowerDB(), func() {}
+				return wtmock.NewTowerDB()
 			},
 		},
 	}
@@ -792,10 +773,7 @@ func TestTowerDB(t *testing.T) {
 
 			for _, test := range tests {
 				t.Run(test.name, func(t *testing.T) {
-					h, cleanup := newTowerDBHarness(
+					h := newTowerDBHarness(t, db.init)
-						t, db.init,
-					)
-					defer cleanup()
 
 					test.run(h)
 				})