From f3ad66b3635e15829b5c0d68e7f213a5c978b71c Mon Sep 17 00:00:00 2001
From: yyforyongyu <yy2452@columbia.edu>
Date: Mon, 17 Oct 2022 09:13:48 +0800
Subject: [PATCH] lntest+lntest: inherit node's public states when restarting
 it

This commit changes how a node's state is created during startup. When
it already has states, only its internal states are reset since these
info are only useful for the finished test. The name `InitRPCClients` is
changed to `Initialize` to reflect the fact that it's more than
initializing RPC clients.
---
 lntemp/node/harness_node.go                   | 27 +++++++++++++------
 lntemp/node/state.go                          | 14 ++++++++++
 .../lnd_rpc_middleware_interceptor_test.go    |  2 +-
 3 files changed, 34 insertions(+), 9 deletions(-)

diff --git a/lntemp/node/harness_node.go b/lntemp/node/harness_node.go
index bdd0fc744..069d23ec8 100644
--- a/lntemp/node/harness_node.go
+++ b/lntemp/node/harness_node.go
@@ -135,15 +135,24 @@ func NewHarnessNode(t *testing.T, cfg *BaseNodeConfig) (*HarnessNode, error) {
 	}, nil
 }
 
-// InitRPCClients initializes a list of RPC clients for the node.
-func (hn *HarnessNode) InitRPCClients(c *grpc.ClientConn) {
+// Initialize creates a list of new RPC clients using the passed connection,
+// initializes the node's internal state and creates a topology watcher.
+func (hn *HarnessNode) Initialize(c *grpc.ClientConn) {
 	hn.conn = c
 
 	// Init all the rpc clients.
 	hn.RPC = rpc.NewHarnessRPC(hn.runCtx, hn.T, c, hn.Name())
 
-	// Init the node's internal state.
-	hn.State = newState(hn.RPC)
+	// Init the node's state.
+	//
+	// If we already have a state, it means we are restarting the node and
+	// we will only reset its internal states. Otherwise we'll create a new
+	// state.
+	if hn.State != nil {
+		hn.State.resetEphermalStates(hn.RPC)
+	} else {
+		hn.State = newState(hn.RPC)
+	}
 
 	// Init the topology watcher.
 	hn.Watcher = newNodeWatcher(hn.RPC, hn.State)
@@ -436,8 +445,9 @@ func (hn *HarnessNode) Start(ctxt context.Context) error {
 		return err
 	}
 
-	// Init all the RPC clients.
-	hn.InitRPCClients(conn)
+	// Init the node by creating the RPC clients, initializing node's
+	// internal state and watcher.
+	hn.Initialize(conn)
 
 	// Wait till the server is starting.
 	if err := hn.WaitUntilStarted(); err != nil {
@@ -479,8 +489,9 @@ func (hn *HarnessNode) InitNode(macBytes []byte) error {
 		}
 	}
 
-	// Init all the RPC clients.
-	hn.InitRPCClients(conn)
+	// Init the node by creating the RPC clients, initializing node's
+	// internal state and watcher.
+	hn.Initialize(conn)
 
 	// Wait till the server is starting.
 	if err := hn.WaitUntilStarted(); err != nil {
diff --git a/lntemp/node/state.go b/lntemp/node/state.go
index e9036c629..e852bf9ab 100644
--- a/lntemp/node/state.go
+++ b/lntemp/node/state.go
@@ -344,3 +344,17 @@ func (s *State) String() string {
 
 	return fmt.Sprintf("\n%s", stateBytes)
 }
+
+// resetEphermalStates resets the current state with a new HarnessRPC and empty
+// private fields which are used to track state only valid for the last test.
+func (s *State) resetEphermalStates(rpc *rpc.HarnessRPC) {
+	s.rpc = rpc
+
+	// Reset ephermal states which are used to record info from finished
+	// tests.
+	s.openChans = &SyncMap[wire.OutPoint, []*OpenChannelUpdate]{}
+	s.closedChans = &SyncMap[wire.OutPoint, *lnrpc.ClosedChannelUpdate]{}
+	s.numChanUpdates = &SyncMap[wire.OutPoint, int]{}
+	s.nodeUpdates = &SyncMap[string, []*lnrpc.NodeUpdate]{}
+	s.policyUpdates = &SyncMap[wire.OutPoint, PolicyUpdate]{}
+}
diff --git a/lntest/itest/lnd_rpc_middleware_interceptor_test.go b/lntest/itest/lnd_rpc_middleware_interceptor_test.go
index f9c1d7b5c..f9a793db4 100644
--- a/lntest/itest/lnd_rpc_middleware_interceptor_test.go
+++ b/lntest/itest/lnd_rpc_middleware_interceptor_test.go
@@ -560,7 +560,7 @@ func middlewareMandatoryTest(ht *lntemp.HarnessTest, node *node.HarnessNode) {
 	// test case. So we need to do the wait and client setup manually here.
 	conn, err := node.ConnectRPC()
 	require.NoError(ht, err)
-	node.InitRPCClients(conn)
+	node.Initialize(conn)
 	err = node.WaitUntilServerActive()
 	require.NoError(ht, err)