htlcswitch: replace defer cleanup with t.Cleanup

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

htlcswitch/mailbox_test.go

@@ -27,7 +27,6 @@ func TestMailBoxCouriers(t *testing.T) {
 	// First, we'll create new instance of the current default mailbox
 	// type.
 	ctx := newMailboxContext(t, time.Now(), testExpiry)
-	defer ctx.mailbox.Stop()
 
 	// We'll be adding 10 message of both types to the mailbox.
 	const numPackets = 10
@@ -215,6 +214,7 @@ func newMailboxContext(t *testing.T, startTime time.Time,
 		expiry:            expiry,
 	})
 	ctx.mailbox.Start()
+	t.Cleanup(ctx.mailbox.Stop)
 
 	return ctx
 }
@@ -311,7 +311,6 @@ func TestMailBoxFailAdd(t *testing.T) {
 		thirdBatchExpiry = thirdBatchStart.Add(expiry)
 	)
 	ctx := newMailboxContext(t, firstBatchStart, expiry)
-	defer ctx.mailbox.Stop()
 
 	failAdds := func(adds []*htlcPacket) {
 		for _, add := range adds {
@@ -377,7 +376,6 @@ func TestMailBoxPacketPrioritization(t *testing.T) {
 	// First, we'll create new instance of the current default mailbox
 	// type.
 	ctx := newMailboxContext(t, time.Now(), testExpiry)
-	defer ctx.mailbox.Stop()
 
 	const numPackets = 5
 
@@ -476,7 +474,6 @@ func TestMailBoxAddExpiry(t *testing.T) {
 	)
 
 	ctx := newMailboxContext(t, firstBatchStart, expiry)
-	defer ctx.mailbox.Stop()
 
 	// Each batch will consist of 10 messages.
 	const numBatchPackets = 10
@@ -503,7 +500,6 @@ func TestMailBoxDuplicateAddPacket(t *testing.T) {
 
 	ctx := newMailboxContext(t, time.Now(), testExpiry)
 	ctx.mailbox.Start()
-	defer ctx.mailbox.Stop()
 
 	addTwice := func(t *testing.T, pkt *htlcPacket) {
 		// The first add should succeed.
@@ -553,7 +549,6 @@ func testMailBoxDust(t *testing.T, chantype channeldb.ChannelType) {
 	t.Parallel()
 
 	ctx := newMailboxContext(t, time.Now(), testExpiry)
-	defer ctx.mailbox.Stop()
 
 	_, _, aliceID, bobID := genIDs()
 

htlcswitch/switch_test.go

@@ -2798,11 +2798,10 @@ func TestLocalPaymentNoForwardingEvents(t *testing.T) {
 	// First, we'll create our traditional three hop network. We'll only be
 	// interacting with and asserting the state of the first end point for
 	// this test.
-	channels, cleanUp, _, err := createClusterChannels(
-		btcutil.SatoshiPerBitcoin*3,
-		btcutil.SatoshiPerBitcoin*5)
+	channels, _, err := createClusterChannels(
+		t, btcutil.SatoshiPerBitcoin*3, btcutil.SatoshiPerBitcoin*5,
+	)
 	require.NoError(t, err, "unable to create channel")
-	defer cleanUp()
 
 	n := newThreeHopNetwork(t, channels.aliceToBob, channels.bobToAlice,
 		channels.bobToCarol, channels.carolToBob, testStartingHeight)
@@ -2856,11 +2855,10 @@ func TestMultiHopPaymentForwardingEvents(t *testing.T) {
 	t.Parallel()
 
 	// First, we'll create our traditional three hop network.
-	channels, cleanUp, _, err := createClusterChannels(
-		btcutil.SatoshiPerBitcoin*3,
-		btcutil.SatoshiPerBitcoin*5)
+	channels, _, err := createClusterChannels(
+		t, btcutil.SatoshiPerBitcoin*3, btcutil.SatoshiPerBitcoin*5,
+	)
 	require.NoError(t, err, "unable to create channel")
-	defer cleanUp()
 
 	n := newThreeHopNetwork(t, channels.aliceToBob, channels.bobToAlice,
 		channels.bobToCarol, channels.carolToBob, testStartingHeight)
@@ -3011,11 +3009,10 @@ func TestUpdateFailMalformedHTLCErrorConversion(t *testing.T) {
 	t.Parallel()
 
 	// First, we'll create our traditional three hop network.
-	channels, cleanUp, _, err := createClusterChannels(
-		btcutil.SatoshiPerBitcoin*3, btcutil.SatoshiPerBitcoin*5,
+	channels, _, err := createClusterChannels(
+		t, btcutil.SatoshiPerBitcoin*3, btcutil.SatoshiPerBitcoin*5,
 	)
 	require.NoError(t, err, "unable to create channel")
-	defer cleanUp()
 
 	n := newThreeHopNetwork(
 		t, channels.aliceToBob, channels.bobToAlice,
@@ -3411,11 +3408,10 @@ func testHtcNotifier(t *testing.T, testOpts []serverOption, iterations int,
 
 	// First, we'll create our traditional three hop
 	// network.
-	channels, cleanUp, _, err := createClusterChannels(
-		btcutil.SatoshiPerBitcoin*3,
-		btcutil.SatoshiPerBitcoin*5)
+	channels, _, err := createClusterChannels(
+		t, btcutil.SatoshiPerBitcoin*3, btcutil.SatoshiPerBitcoin*5,
+	)
 	require.NoError(t, err, "unable to create channel")
-	defer cleanUp()
 
 	// Mock time so that all events are reported with a static timestamp.
 	now := time.Now()
@@ -3429,31 +3425,31 @@ func testHtcNotifier(t *testing.T, testOpts []serverOption, iterations int,
 	if err := aliceNotifier.Start(); err != nil {
 		t.Fatalf("could not start alice notifier")
 	}
-	defer func() {
+	t.Cleanup(func() {
 		if err := aliceNotifier.Stop(); err != nil {
 			t.Fatalf("failed to stop alice notifier: %v", err)
 		}
-	}()
+	})
 
 	bobNotifier := NewHtlcNotifier(mockTime)
 	if err := bobNotifier.Start(); err != nil {
 		t.Fatalf("could not start bob notifier")
 	}
-	defer func() {
+	t.Cleanup(func() {
 		if err := bobNotifier.Stop(); err != nil {
 			t.Fatalf("failed to stop bob notifier: %v", err)
 		}
-	}()
+	})
 
 	carolNotifier := NewHtlcNotifier(mockTime)
 	if err := carolNotifier.Start(); err != nil {
 		t.Fatalf("could not start carol notifier")
 	}
-	defer func() {
+	t.Cleanup(func() {
 		if err := carolNotifier.Stop(); err != nil {
 			t.Fatalf("failed to stop carol notifier: %v", err)
 		}
-	}()
+	})
 
 	// Create a notifier server option which will set our htlc notifiers
 	// for the three hop network.
@@ -3475,7 +3471,7 @@ func testHtcNotifier(t *testing.T, testOpts []serverOption, iterations int,
 		t.Fatalf("unable to start three hop "+
 			"network: %v", err)
 	}
-	defer n.stop()
+	t.Cleanup(n.stop)
 
 	// Before we forward anything, subscribe to htlc events
 	// from each notifier.
@@ -3484,21 +3480,21 @@ func testHtcNotifier(t *testing.T, testOpts []serverOption, iterations int,
 		t.Fatalf("could not subscribe to alice's"+
 			" events: %v", err)
 	}
-	defer aliceEvents.Cancel()
+	t.Cleanup(aliceEvents.Cancel)
 
 	bobEvents, err := bobNotifier.SubscribeHtlcEvents()
 	if err != nil {
 		t.Fatalf("could not subscribe to bob's"+
 			" events: %v", err)
 	}
-	defer bobEvents.Cancel()
+	t.Cleanup(bobEvents.Cancel)
 
 	carolEvents, err := carolNotifier.SubscribeHtlcEvents()
 	if err != nil {
 		t.Fatalf("could not subscribe to carol's"+
 			" events: %v", err)
 	}
-	defer carolEvents.Cancel()
+	t.Cleanup(carolEvents.Cancel)
 
 	// Send multiple payments, as specified by the test to test incrementing
 	// of htlc ids.
@@ -4118,11 +4114,10 @@ func TestSwitchDustForwarding(t *testing.T) {
 	// - Bob has a dust limit of 800sats with Alice
 	// - Bob has a dust limit of 200sats with Carol
 	// - Carol has a dust limit of 800sats with Bob
-	channels, cleanUp, _, err := createClusterChannels(
-		btcutil.SatoshiPerBitcoin, btcutil.SatoshiPerBitcoin,
+	channels, _, err := createClusterChannels(
+		t, btcutil.SatoshiPerBitcoin, btcutil.SatoshiPerBitcoin,
 	)
 	require.NoError(t, err)
-	defer cleanUp()
 
 	n := newThreeHopNetwork(
 		t, channels.aliceToBob, channels.bobToAlice,
@@ -5114,10 +5109,9 @@ func testSwitchHandlePacketForward(t *testing.T, zeroConf, private,
 	t.Parallel()
 
 	// Create a link for Alice that we'll add to the switch.
-	aliceLink, _, _, _, cleanUp, _, err :=
-		newSingleLinkTestHarness(btcutil.SatoshiPerBitcoin, 0)
+	aliceLink, _, _, _, _, err :=
+		newSingleLinkTestHarness(t, btcutil.SatoshiPerBitcoin, 0)
 	require.NoError(t, err)
-	defer cleanUp()
 
 	s, err := initSwitchWithTempDB(t, testStartingHeight)
 	if err != nil {

htlcswitch/test_utils.go

@@ -7,7 +7,6 @@ import (
 	"encoding/binary"
 	"encoding/hex"
 	"fmt"
-	"io/ioutil"
 	"net"
 	"os"
 	"runtime"
@@ -124,10 +123,10 @@ type testLightningChannel struct {
 // representations.
 //
 // TODO(roasbeef): need to factor out, similar func re-used in many parts of codebase
-func createTestChannel(alicePrivKey, bobPrivKey []byte,
+func createTestChannel(t *testing.T, alicePrivKey, bobPrivKey []byte,
 	aliceAmount, bobAmount, aliceReserve, bobReserve btcutil.Amount,
 	chanID lnwire.ShortChannelID) (*testLightningChannel,
-	*testLightningChannel, func(), error) {
+	*testLightningChannel, error) {
 
 	aliceKeyPriv, aliceKeyPub := btcec.PrivKeyFromBytes(alicePrivKey)
 	bobKeyPriv, bobKeyPub := btcec.PrivKeyFromBytes(bobPrivKey)
@@ -160,7 +159,7 @@ func createTestChannel(alicePrivKey, bobPrivKey []byte,
 	var hash [sha256.Size]byte
 	randomSeed, err := generateRandomBytes(sha256.Size)
 	if err != nil {
-		return nil, nil, nil, err
+		return nil, nil, err
 	}
 	copy(hash[:], randomSeed)
 
@@ -209,23 +208,23 @@ func createTestChannel(alicePrivKey, bobPrivKey []byte,
 
 	bobRoot, err := chainhash.NewHash(bobKeyPriv.Serialize())
 	if err != nil {
-		return nil, nil, nil, err
+		return nil, nil, err
 	}
 	bobPreimageProducer := shachain.NewRevocationProducer(*bobRoot)
 	bobFirstRevoke, err := bobPreimageProducer.AtIndex(0)
 	if err != nil {
-		return nil, nil, nil, err
+		return nil, nil, err
 	}
 	bobCommitPoint := input.ComputeCommitmentPoint(bobFirstRevoke[:])
 
 	aliceRoot, err := chainhash.NewHash(aliceKeyPriv.Serialize())
 	if err != nil {
-		return nil, nil, nil, err
+		return nil, nil, err
 	}
 	alicePreimageProducer := shachain.NewRevocationProducer(*aliceRoot)
 	aliceFirstRevoke, err := alicePreimageProducer.AtIndex(0)
 	if err != nil {
-		return nil, nil, nil, err
+		return nil, nil, err
 	}
 	aliceCommitPoint := input.ComputeCommitmentPoint(aliceFirstRevoke[:])
 
@@ -235,33 +234,29 @@ func createTestChannel(alicePrivKey, bobPrivKey []byte,
 		isAliceInitiator, 0,
 	)
 	if err != nil {
-		return nil, nil, nil, err
+		return nil, nil, err
 	}
 
-	alicePath, err := ioutil.TempDir("", "alicedb")
+	dbAlice, err := channeldb.Open(t.TempDir())
 	if err != nil {
-		return nil, nil, nil, err
+		return nil, nil, err
 	}
+	t.Cleanup(func() {
+		require.NoError(t, dbAlice.Close())
+	})
 
-	dbAlice, err := channeldb.Open(alicePath)
+	dbBob, err := channeldb.Open(t.TempDir())
 	if err != nil {
-		return nil, nil, nil, err
-	}
-
-	bobPath, err := ioutil.TempDir("", "bobdb")
-	if err != nil {
-		return nil, nil, nil, err
-	}
-
-	dbBob, err := channeldb.Open(bobPath)
-	if err != nil {
-		return nil, nil, nil, err
+		return nil, nil, err
 	}
+	t.Cleanup(func() {
+		require.NoError(t, dbBob.Close())
+	})
 
 	estimator := chainfee.NewStaticEstimator(6000, 0)
 	feePerKw, err := estimator.EstimateFeePerKW(1)
 	if err != nil {
-		return nil, nil, nil, err
+		return nil, nil, err
 	}
 	commitFee := feePerKw.FeeForWeight(724)
 
@@ -333,18 +328,11 @@ func createTestChannel(alicePrivKey, bobPrivKey []byte,
 	}
 
 	if err := aliceChannelState.SyncPending(bobAddr, broadcastHeight); err != nil {
-		return nil, nil, nil, err
+		return nil, nil, err
 	}
 
 	if err := bobChannelState.SyncPending(aliceAddr, broadcastHeight); err != nil {
-		return nil, nil, nil, err
-	}
-
-	cleanUpFunc := func() {
-		dbAlice.Close()
-		dbBob.Close()
-		os.RemoveAll(bobPath)
-		os.RemoveAll(alicePath)
+		return nil, nil, err
 	}
 
 	aliceSigner := &mock.SingleSigner{Privkey: aliceKeyPriv}
@@ -355,7 +343,7 @@ func createTestChannel(alicePrivKey, bobPrivKey []byte,
 		aliceSigner, aliceChannelState, alicePool,
 	)
 	if err != nil {
-		return nil, nil, nil, err
+		return nil, nil, err
 	}
 	alicePool.Start()
 
@@ -364,7 +352,7 @@ func createTestChannel(alicePrivKey, bobPrivKey []byte,
 		bobSigner, bobChannelState, bobPool,
 	)
 	if err != nil {
-		return nil, nil, nil, err
+		return nil, nil, err
 	}
 	bobPool.Start()
 
@@ -372,18 +360,18 @@ func createTestChannel(alicePrivKey, bobPrivKey []byte,
 	// having Alice and Bob extend their revocation windows to each other.
 	aliceNextRevoke, err := channelAlice.NextRevocationKey()
 	if err != nil {
-		return nil, nil, nil, err
+		return nil, nil, err
 	}
 	if err := channelBob.InitNextRevocation(aliceNextRevoke); err != nil {
-		return nil, nil, nil, err
+		return nil, nil, err
 	}
 
 	bobNextRevoke, err := channelBob.NextRevocationKey()
 	if err != nil {
-		return nil, nil, nil, err
+		return nil, nil, err
 	}
 	if err := channelAlice.InitNextRevocation(bobNextRevoke); err != nil {
-		return nil, nil, nil, err
+		return nil, nil, err
 	}
 
 	restoreAlice := func() (*lnwallet.LightningChannel, error) {
@@ -487,8 +475,7 @@ func createTestChannel(alicePrivKey, bobPrivKey []byte,
 		restore: restoreBob,
 	}
 
-	return testLightningChannelAlice, testLightningChannelBob, cleanUpFunc,
-		nil
+	return testLightningChannelAlice, testLightningChannelBob, nil
 }
 
 // getChanID retrieves the channel point from an lnnwire message.
@@ -872,34 +859,28 @@ type clusterChannels struct {
 
 // createClusterChannels creates lightning channels which are needed for
 // network cluster to be initialized.
-func createClusterChannels(aliceToBob, bobToCarol btcutil.Amount) (
-	*clusterChannels, func(), func() (*clusterChannels, error), error) {
+func createClusterChannels(t *testing.T, aliceToBob, bobToCarol btcutil.Amount) (
+	*clusterChannels, func() (*clusterChannels, error), error) {
 
 	_, _, firstChanID, secondChanID := genIDs()
 
 	// Create lightning channels between Alice<->Bob and Bob<->Carol
-	aliceChannel, firstBobChannel, cleanAliceBob, err :=
-		createTestChannel(alicePrivKey, bobPrivKey, aliceToBob,
-			aliceToBob, 0, 0, firstChanID)
+	aliceChannel, firstBobChannel, err := createTestChannel(t, alicePrivKey,
+		bobPrivKey, aliceToBob, aliceToBob, 0, 0, firstChanID,
+	)
 	if err != nil {
-		return nil, nil, nil, errors.Errorf("unable to create "+
+		return nil, nil, errors.Errorf("unable to create "+
 			"alice<->bob channel: %v", err)
 	}
 
-	secondBobChannel, carolChannel, cleanBobCarol, err :=
-		createTestChannel(bobPrivKey, carolPrivKey, bobToCarol,
-			bobToCarol, 0, 0, secondChanID)
+	secondBobChannel, carolChannel, err := createTestChannel(t, bobPrivKey,
+		carolPrivKey, bobToCarol, bobToCarol, 0, 0, secondChanID,
+	)
 	if err != nil {
-		cleanAliceBob()
-		return nil, nil, nil, errors.Errorf("unable to create "+
+		return nil, nil, errors.Errorf("unable to create "+
 			"bob<->carol channel: %v", err)
 	}
 
-	cleanUp := func() {
-		cleanAliceBob()
-		cleanBobCarol()
-	}
-
 	restoreFromDb := func() (*clusterChannels, error) {
 
 		a2b, err := aliceChannel.restore()
@@ -935,7 +916,7 @@ func createClusterChannels(aliceToBob, bobToCarol btcutil.Amount) (
 		bobToAlice: firstBobChannel.channel,
 		bobToCarol: secondBobChannel.channel,
 		carolToBob: carolChannel.channel,
-	}, cleanUp, restoreFromDb, nil
+	}, restoreFromDb, nil
 }
 
 // newThreeHopNetwork function creates the following topology and returns the
@@ -1066,22 +1047,22 @@ func serverOptionRejectHtlc(alice, bob, carol bool) serverOption {
 
 // createTwoClusterChannels creates lightning channels which are needed for
 // a 2 hop network cluster to be initialized.
-func createTwoClusterChannels(aliceToBob, bobToCarol btcutil.Amount) (
-	*testLightningChannel, *testLightningChannel,
-	func(), error) {
+func createTwoClusterChannels(t *testing.T, aliceToBob,
+	bobToCarol btcutil.Amount) (*testLightningChannel,
+	*testLightningChannel, error) {
 
 	_, _, firstChanID, _ := genIDs()
 
 	// Create lightning channels between Alice<->Bob and Bob<->Carol
-	alice, bob, cleanAliceBob, err :=
-		createTestChannel(alicePrivKey, bobPrivKey, aliceToBob,
-			aliceToBob, 0, 0, firstChanID)
+	alice, bob, err := createTestChannel(t, alicePrivKey, bobPrivKey,
+		aliceToBob, aliceToBob, 0, 0, firstChanID,
+	)
 	if err != nil {
-		return nil, nil, nil, errors.Errorf("unable to create "+
+		return nil, nil, errors.Errorf("unable to create "+
 			"alice<->bob channel: %v", err)
 	}
 
-	return alice, bob, cleanAliceBob, nil
+	return alice, bob, nil
 }
 
 // hopNetwork is the base struct for two and three hop networks
@@ -1211,8 +1192,8 @@ type twoHopNetwork struct {
 	bobChannelLink *channelLink
 }
 
-// newTwoHopNetwork function creates the following topology and returns the
-// control object to manage this cluster:
+// newTwoHopNetwork function creates and starts the following topology and
+// returns the control object to manage this cluster:
 //
 // alice                      bob
 // server - <-connection-> - server
@@ -1265,7 +1246,7 @@ func newTwoHopNetwork(t testing.TB,
 		t.Fatal(err)
 	}
 
-	return &twoHopNetwork{
+	n := &twoHopNetwork{
 		aliceServer:      aliceServer,
 		aliceChannelLink: aliceChannelLink.(*channelLink),
 
@@ -1274,6 +1255,11 @@ func newTwoHopNetwork(t testing.TB,
 
 		hopNetwork: *hopNetwork,
 	}
+
+	require.NoError(t, n.start())
+	t.Cleanup(n.stop)
+
+	return n
 }
 
 // start starts the two hop network alice,bob servers.
@@ -1398,7 +1384,7 @@ func waitLinksEligible(links map[string]*channelLink) error {
 }
 
 // timeout implements a test level timeout.
-func timeout(t *testing.T) func() {
+func timeout() func() {
 	done := make(chan struct{})
 	go func() {
 		select {