From 66dae6ecf7e82213daa08f05506715daf02e31c0 Mon Sep 17 00:00:00 2001
From: yyforyongyu <yy2452@columbia.edu>
Date: Wed, 25 Aug 2021 15:22:49 +0800
Subject: [PATCH] itest: put node.CloseChannel inside wait

---
 lntest/harness.go                   | 60 ++++++++++++++---------------
 lntest/itest/lnd_revocation_test.go | 39 +++++--------------
 2 files changed, 37 insertions(+), 62 deletions(-)

diff --git a/lntest/harness.go b/lntest/harness.go
index d2648a5d3..8e5fe8786 100644
--- a/lntest/harness.go
+++ b/lntest/harness.go
@@ -1263,56 +1263,52 @@ func (n *NetworkHarness) CloseChannel(lnNode *HarnessNode,
 		}
 	}
 
-	closeReq := &lnrpc.CloseChannelRequest{
-		ChannelPoint: cp,
-		Force:        force,
-	}
-	closeRespStream, err := lnNode.CloseChannel(ctx, closeReq)
-	if err != nil {
-		return nil, nil, fmt.Errorf("unable to close channel: %v", err)
-	}
+	var (
+		closeRespStream lnrpc.Lightning_CloseChannelClient
+		closeTxid       *chainhash.Hash
+	)
 
-	errChan := make(chan error)
-	fin := make(chan *chainhash.Hash)
-	go func() {
-		// Consume the "channel close" update in order to wait for the closing
-		// transaction to be broadcast, then wait for the closing tx to be seen
-		// within the network.
+	err = wait.NoError(func() error {
+		closeReq := &lnrpc.CloseChannelRequest{
+			ChannelPoint: cp, Force: force,
+		}
+		closeRespStream, err = lnNode.CloseChannel(ctx, closeReq)
+		if err != nil {
+			return fmt.Errorf("unable to close channel: %v", err)
+		}
+
+		// Consume the "channel close" update in order to wait for the
+		// closing transaction to be broadcast, then wait for the
+		// closing tx to be seen within the network.
 		closeResp, err := closeRespStream.Recv()
 		if err != nil {
-			errChan <- fmt.Errorf("unable to recv() from close "+
+			return fmt.Errorf("unable to recv() from close "+
 				"stream: %v", err)
-			return
 		}
 		pendingClose, ok := closeResp.Update.(*lnrpc.CloseStatusUpdate_ClosePending)
 		if !ok {
-			errChan <- fmt.Errorf("expected channel close update, "+
+			return fmt.Errorf("expected channel close update, "+
 				"instead got %v", pendingClose)
-			return
 		}
 
-		closeTxid, err := chainhash.NewHash(pendingClose.ClosePending.Txid)
+		closeTxid, err = chainhash.NewHash(
+			pendingClose.ClosePending.Txid,
+		)
 		if err != nil {
-			errChan <- fmt.Errorf("unable to decode closeTxid: "+
+			return fmt.Errorf("unable to decode closeTxid: "+
 				"%v", err)
-			return
 		}
 		if err := n.waitForTxInMempool(ctx, *closeTxid); err != nil {
-			errChan <- fmt.Errorf("error while waiting for "+
+			return fmt.Errorf("error while waiting for "+
 				"broadcast tx: %v", err)
-			return
 		}
-		fin <- closeTxid
-	}()
-
-	// Wait until either the deadline for the context expires, an error
-	// occurs, or the channel close update is received.
-	select {
-	case err := <-errChan:
+		return nil
+	}, ChannelCloseTimeout)
+	if err != nil {
 		return nil, nil, err
-	case closeTxid := <-fin:
-		return closeRespStream, closeTxid, nil
 	}
+
+	return closeRespStream, closeTxid, nil
 }
 
 // WaitForChannelClose waits for a notification from the passed channel close
diff --git a/lntest/itest/lnd_revocation_test.go b/lntest/itest/lnd_revocation_test.go
index 97d6c3eb1..d26356b83 100644
--- a/lntest/itest/lnd_revocation_test.go
+++ b/lntest/itest/lnd_revocation_test.go
@@ -17,6 +17,7 @@ import (
 	"github.com/lightningnetwork/lnd/lnrpc/wtclientrpc"
 	"github.com/lightningnetwork/lnd/lntest"
 	"github.com/lightningnetwork/lnd/lntest/wait"
+	"github.com/stretchr/testify/require"
 )
 
 // testRevokedCloseRetribution tests that Carol is able carry out
@@ -159,22 +160,11 @@ func testRevokedCloseRetribution(net *lntest.NetworkHarness, t *harnessTest) {
 	// broadcasting his current channel state. This is actually the
 	// commitment transaction of a prior *revoked* state, so he'll soon
 	// feel the wrath of Carol's retribution.
-	var closeUpdates lnrpc.Lightning_CloseChannelClient
 	force := true
-	err = wait.Predicate(func() bool {
-		closeUpdates, _, err = net.CloseChannel(
-			net.Bob, chanPoint, force,
-		)
-		if err != nil {
-			predErr = err
-			return false
-		}
-
-		return true
-	}, defaultTimeout)
-	if err != nil {
-		t.Fatalf("unable to close channel: %v", predErr)
-	}
+	closeUpdates, _, err := net.CloseChannel(
+		net.Bob, chanPoint, force,
+	)
+	require.NoError(t.t, err, "unable to close channel")
 
 	// Wait for Bob's breach transaction to show up in the mempool to ensure
 	// that Carol's node has started waiting for confirmations.
@@ -386,22 +376,11 @@ func testRevokedCloseRetributionZeroValueRemoteOutput(net *lntest.NetworkHarness
 	// broadcasting her current channel state. This is actually the
 	// commitment transaction of a prior *revoked* state, so she'll soon
 	// feel the wrath of Dave's retribution.
-	var (
-		closeUpdates lnrpc.Lightning_CloseChannelClient
-		closeTxID    *chainhash.Hash
-		closeErr     error
-	)
-
 	force := true
-	err = wait.Predicate(func() bool {
-		closeUpdates, closeTxID, closeErr = net.CloseChannel(
-			carol, chanPoint, force,
-		)
-		return closeErr == nil
-	}, defaultTimeout)
-	if err != nil {
-		t.Fatalf("unable to close channel: %v", closeErr)
-	}
+	closeUpdates, closeTxID, closeErr := net.CloseChannel(
+		carol, chanPoint, force,
+	)
+	require.NoError(t.t, closeErr, "unable to close channel")
 
 	// Query the mempool for the breaching closing transaction, this should
 	// be broadcast by Carol when she force closes the channel above.