From 9dfbdfd29e65877bc1b087833efc47e0b768d1f6 Mon Sep 17 00:00:00 2001
From: yyforyongyu <yy2452@columbia.edu>
Date: Tue, 14 Mar 2023 21:24:50 +0800
Subject: [PATCH] itest+lntest: make sure states are cleaned when tests end

This commit changes how the node's state is updated to make sure the
test cleans up the node's state.

Also `testLookupHtlcResolution` is fixed with a cleanup.
---
 itest/lnd_htlc_test.go |  1 +
 lntest/harness.go      | 61 +++++++++++++++++++++++++++---------------
 2 files changed, 41 insertions(+), 21 deletions(-)

diff --git a/itest/lnd_htlc_test.go b/itest/lnd_htlc_test.go
index 4e73a385b6..fc04fe8d0f 100644
--- a/itest/lnd_htlc_test.go
+++ b/itest/lnd_htlc_test.go
@@ -24,6 +24,7 @@ func testLookupHtlcResolution(ht *lntest.HarnessTest) {
 	cp := ht.OpenChannel(
 		alice, carol, lntest.OpenChannelParams{Amt: chanAmt},
 	)
+	defer ht.CloseChannel(alice, cp)
 
 	// Channel should be ready for payments.
 	const payAmt = 100
diff --git a/lntest/harness.go b/lntest/harness.go
index 95cc1afbbc..819a0a8d3e 100644
--- a/lntest/harness.go
+++ b/lntest/harness.go
@@ -299,9 +299,6 @@ func (h *HarnessTest) resetStandbyNodes(t *testing.T) {
 		// config for the coming test. This will also inherit the
 		// test's running context.
 		h.RestartNodeWithExtraArgs(hn, hn.Cfg.OriginalExtraArgs)
-
-		// Update the node's internal state.
-		hn.UpdateState()
 	}
 }
 
@@ -433,8 +430,19 @@ func (h *HarnessTest) cleanupStandbyNode(hn *node.HarnessNode) {
 	// Delete all payments made from this test.
 	hn.RPC.DeleteAllPayments()
 
-	// Finally, check the node is in a clean state for the following tests.
-	h.validateNodeState(hn)
+	// Check the node's current state with timeout.
+	//
+	// NOTE: we need to do this in a `wait` because it takes some time for
+	// the node to update its internal state. Once the RPCs are synced we
+	// can then remove this wait.
+	err := wait.NoError(func() error {
+		// Update the node's internal state.
+		hn.UpdateState()
+
+		// Check the node is in a clean state for the following tests.
+		return h.validateNodeState(hn)
+	}, wait.DefaultTimeout)
+	require.NoError(h, err, "timeout checking node's state")
 }
 
 // removeConnectionns will remove all connections made on the standby nodes
@@ -752,32 +760,43 @@ func (h *HarnessTest) SetFeeEstimateWithConf(
 
 // validateNodeState checks that the node doesn't have any uncleaned states
 // which will affect its following tests.
-func (h *HarnessTest) validateNodeState(hn *node.HarnessNode) {
-	errStr := func(subject string) string {
-		return fmt.Sprintf("%s: found %s channels, please close "+
+func (h *HarnessTest) validateNodeState(hn *node.HarnessNode) error {
+	errStr := func(subject string) error {
+		return fmt.Errorf("%s: found %s channels, please close "+
 			"them properly", hn.Name(), subject)
 	}
 	// If the node still has open channels, it's most likely that the
 	// current test didn't close it properly.
-	require.Zerof(h, hn.State.OpenChannel.Active, errStr("active"))
-	require.Zerof(h, hn.State.OpenChannel.Public, errStr("public"))
-	require.Zerof(h, hn.State.OpenChannel.Private, errStr("private"))
-	require.Zerof(h, hn.State.OpenChannel.Pending, errStr("pending open"))
+	if hn.State.OpenChannel.Active != 0 {
+		return errStr("active")
+	}
+	if hn.State.OpenChannel.Public != 0 {
+		return errStr("public")
+	}
+	if hn.State.OpenChannel.Private != 0 {
+		return errStr("private")
+	}
+	if hn.State.OpenChannel.Pending != 0 {
+		return errStr("pending open")
+	}
 
 	// The number of pending force close channels should be zero.
-	require.Zerof(h, hn.State.CloseChannel.PendingForceClose,
-		errStr("pending force"))
+	if hn.State.CloseChannel.PendingForceClose != 0 {
+		return errStr("pending force")
+	}
 
 	// The number of waiting close channels should be zero.
-	require.Zerof(h, hn.State.CloseChannel.WaitingClose,
-		errStr("waiting close"))
+	if hn.State.CloseChannel.WaitingClose != 0 {
+		return errStr("waiting close")
+	}
 
 	// Ths number of payments should be zero.
-	// TODO(yy): no need to check since it's deleted in the cleanup? Or
-	// check it in a wait?
-	require.Zerof(h, hn.State.Payment.Total, "%s: found "+
-		"uncleaned payments, please delete all of them properly",
-		hn.Name())
+	if hn.State.Payment.Total != 0 {
+		return fmt.Errorf("%s: found uncleaned payments, please "+
+			"delete all of them properly", hn.Name())
+	}
+
+	return nil
 }
 
 // GetChanPointFundingTxid takes a channel point and converts it into a chain