macarooons: use T.TempDir to create temporary test directory

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

macaroons/service_test.go

@@ -3,8 +3,6 @@ package macaroons_test
 import (
 	"context"
 	"encoding/hex"
-	"io/ioutil"
-	"os"
 	"path"
 	"testing"
 
@@ -33,11 +31,9 @@ var (
 // default password of 'hello'. Only the path to the temporary
 // DB file is returned, because the service will open the file
 // and read the store on its own.
-func setupTestRootKeyStorage(t *testing.T) (string, kvdb.Backend) {
-	tempDir, err := ioutil.TempDir("", "macaroonstore-")
-	require.NoError(t, err, "Error creating temp dir")
+func setupTestRootKeyStorage(t *testing.T) kvdb.Backend {
 	db, err := kvdb.Create(
-		kvdb.BoltBackendName, path.Join(tempDir, "macaroons.db"), true,
+		kvdb.BoltBackendName, path.Join(t.TempDir(), "macaroons.db"), true,
 		kvdb.DefaultDBTimeout,
 	)
 	require.NoError(t, err, "Error opening store DB")
@@ -49,15 +45,14 @@ func setupTestRootKeyStorage(t *testing.T) (string, kvdb.Backend) {
 	defer store.Close()
 	err = store.CreateUnlock(&defaultPw)
 	require.NoError(t, err, "error creating unlock")
-	return tempDir, db
+	return db
 }
 
 // TestNewService tests the creation of the macaroon service.
 func TestNewService(t *testing.T) {
 	// First, initialize a dummy DB file with a store that the service
 	// can read from. Make sure the file is removed in the end.
-	tempDir, db := setupTestRootKeyStorage(t)
-	defer os.RemoveAll(tempDir)
+	db := setupTestRootKeyStorage(t)
 
 	rootKeyStore, err := macaroons.NewRootKeyStorage(db)
 	require.NoError(t, err)
@@ -107,8 +102,7 @@ func TestNewService(t *testing.T) {
 // incoming context.
 func TestValidateMacaroon(t *testing.T) {
 	// First, initialize the service and unlock it.
-	tempDir, db := setupTestRootKeyStorage(t)
-	defer os.RemoveAll(tempDir)
+	db := setupTestRootKeyStorage(t)
 	rootKeyStore, err := macaroons.NewRootKeyStorage(db)
 	require.NoError(t, err)
 	service, err := macaroons.NewService(
@@ -154,8 +148,7 @@ func TestValidateMacaroon(t *testing.T) {
 func TestListMacaroonIDs(t *testing.T) {
 	// First, initialize a dummy DB file with a store that the service
 	// can read from. Make sure the file is removed in the end.
-	tempDir, db := setupTestRootKeyStorage(t)
-	defer os.RemoveAll(tempDir)
+	db := setupTestRootKeyStorage(t)
 
 	// Second, create the new service instance, unlock it and pass in a
 	// checker that we expect it to add to the bakery.
@@ -188,8 +181,7 @@ func TestDeleteMacaroonID(t *testing.T) {
 
 	// First, initialize a dummy DB file with a store that the service
 	// can read from. Make sure the file is removed in the end.
-	tempDir, db := setupTestRootKeyStorage(t)
-	defer os.RemoveAll(tempDir)
+	db := setupTestRootKeyStorage(t)
 
 	// Second, create the new service instance, unlock it and pass in a
 	// checker that we expect it to add to the bakery.

macaroons/store_test.go

@@ -3,8 +3,6 @@ package macaroons_test
 import (
 	"context"
 	"crypto/rand"
-	"io/ioutil"
-	"os"
 	"path"
 	"testing"
 
@@ -22,24 +20,17 @@ var (
 
 // newTestStore creates a new bolt DB in a temporary directory and then
 // initializes a root key storage for that DB.
-func newTestStore(t *testing.T) (string, func(), *macaroons.RootKeyStorage) {
-	tempDir, err := ioutil.TempDir("", "macaroonstore-")
-	require.NoError(t, err)
+func newTestStore(t *testing.T) (string, *macaroons.RootKeyStorage) {
+	tempDir := t.TempDir()
 
-	cleanup, store := openTestStore(t, tempDir)
-	cleanup2 := func() {
-		cleanup()
-		_ = os.RemoveAll(tempDir)
-	}
+	store := openTestStore(t, tempDir)
 
-	return tempDir, cleanup2, store
+	return tempDir, store
 }
 
 // openTestStore opens an existing bolt DB and then initializes a root key
 // storage for that DB.
-func openTestStore(t *testing.T, tempDir string) (func(),
-	*macaroons.RootKeyStorage) {
-
+func openTestStore(t *testing.T, tempDir string) *macaroons.RootKeyStorage {
 	db, err := kvdb.Create(
 		kvdb.BoltBackendName, path.Join(tempDir, "weks.db"), true,
 		kvdb.DefaultDBTimeout,
@@ -52,19 +43,18 @@ func openTestStore(t *testing.T, tempDir string) (func(),
 		t.Fatalf("Error creating root key store: %v", err)
 	}
 
-	cleanup := func() {
+	t.Cleanup(func() {
 		_ = store.Close()
 		_ = db.Close()
-	}
+	})
 
-	return cleanup, store
+	return store
 }
 
 // TestStore tests the normal use cases of the store like creating, unlocking,
 // reading keys and closing it.
 func TestStore(t *testing.T) {
-	tempDir, cleanup, store := newTestStore(t)
-	defer cleanup()
+	tempDir, store := newTestStore(t)
 
 	_, _, err := store.RootKey(context.TODO())
 	require.Equal(t, macaroons.ErrStoreLocked, err)
@@ -114,7 +104,7 @@ func TestStore(t *testing.T) {
 	// Between here and the re-opening of the store, it's possible to get
 	// a double-close, but that's not such a big deal since the tests will
 	// fail anyway in that case.
-	_, store = openTestStore(t, tempDir)
+	store = openTestStore(t, tempDir)
 
 	err = store.CreateUnlock(&badpw)
 	require.Equal(t, snacl.ErrInvalidPassword, err)
@@ -144,8 +134,7 @@ func TestStore(t *testing.T) {
 // TestStoreGenerateNewRootKey tests that a root key can be replaced with a new
 // one in the store without changing the password.
 func TestStoreGenerateNewRootKey(t *testing.T) {
-	_, cleanup, store := newTestStore(t)
-	defer cleanup()
+	_, store := newTestStore(t)
 
 	// The store must be unlocked to replace the root key.
 	err := store.GenerateNewRootKey()
@@ -172,8 +161,7 @@ func TestStoreGenerateNewRootKey(t *testing.T) {
 
 // TestStoreSetRootKey tests that a root key can be set to a specified value.
 func TestStoreSetRootKey(t *testing.T) {
-	_, cleanup, store := newTestStore(t)
-	defer cleanup()
+	_, store := newTestStore(t)
 
 	// Create a new random key
 	rootKey := make([]byte, 32)
@@ -209,8 +197,7 @@ func TestStoreSetRootKey(t *testing.T) {
 // TestStoreChangePassword tests that the password for the store can be changed
 // without changing the root key.
 func TestStoreChangePassword(t *testing.T) {
-	tempDir, cleanup, store := newTestStore(t)
-	defer cleanup()
+	tempDir, store := newTestStore(t)
 
 	// The store must be unlocked to replace the root key.
 	err := store.ChangePassword(nil, nil)
@@ -251,7 +238,7 @@ func TestStoreChangePassword(t *testing.T) {
 	require.Error(t, err)
 
 	// Let's open it again and try unlocking with the new password.
-	_, store = openTestStore(t, tempDir)
+	store = openTestStore(t, tempDir)
 	err = store.CreateUnlock(&newPw)
 	require.NoError(t, err)