From 0fac6c400ec51bac612030cb75d409d543b2a698 Mon Sep 17 00:00:00 2001
From: Oliver Gugger <gugger@gmail.com>
Date: Wed, 6 Oct 2021 14:12:11 +0200
Subject: [PATCH] itest: wait for node to fully start up

---
 .../lnd_rpc_middleware_interceptor_test.go    |  4 +++-
 lntest/node.go                                | 23 ++++++++++++++++++-
 2 files changed, 25 insertions(+), 2 deletions(-)

diff --git a/lntest/itest/lnd_rpc_middleware_interceptor_test.go b/lntest/itest/lnd_rpc_middleware_interceptor_test.go
index e2381a40d..54e0b8db6 100644
--- a/lntest/itest/lnd_rpc_middleware_interceptor_test.go
+++ b/lntest/itest/lnd_rpc_middleware_interceptor_test.go
@@ -403,7 +403,9 @@ func middlewareMandatoryTest(t *testing.T, node *lntest.HarnessNode,
 	// test case. So we need to do the wait and client setup manually here.
 	conn, err := node.ConnectRPC(true)
 	require.NoError(t, err)
-	err = node.WaitUntilStarted(conn, defaultTimeout)
+	err = node.WaitUntilStateReached(
+		conn, defaultTimeout, lnrpc.WalletState_RPC_ACTIVE,
+	)
 	require.NoError(t, err)
 	node.LightningClient = lnrpc.NewLightningClient(conn)
 
diff --git a/lntest/node.go b/lntest/node.go
index 2957e2c65..f164ad070 100644
--- a/lntest/node.go
+++ b/lntest/node.go
@@ -887,6 +887,27 @@ func (hn *HarnessNode) start(lndBinary string, lndError chan<- error,
 func (hn *HarnessNode) WaitUntilStarted(conn grpc.ClientConnInterface,
 	timeout time.Duration) error {
 
+	return hn.waitForState(conn, timeout, func(s lnrpc.WalletState) bool {
+		return s != lnrpc.WalletState_WAITING_TO_START
+	})
+}
+
+// WaitUntilStateReached waits until the given wallet state (or one of the
+// states following it) has been reached.
+func (hn *HarnessNode) WaitUntilStateReached(conn grpc.ClientConnInterface,
+	timeout time.Duration, desiredState lnrpc.WalletState) error {
+
+	return hn.waitForState(conn, timeout, func(s lnrpc.WalletState) bool {
+		return s >= desiredState
+	})
+}
+
+// waitForState waits until the current node state fulfills the given
+// predicate.
+func (hn *HarnessNode) waitForState(conn grpc.ClientConnInterface,
+	timeout time.Duration,
+	predicate func(state lnrpc.WalletState) bool) error {
+
 	stateClient := lnrpc.NewStateClient(conn)
 	ctx, cancel := context.WithCancel(context.Background())
 	defer cancel()
@@ -908,7 +929,7 @@ func (hn *HarnessNode) WaitUntilStarted(conn grpc.ClientConnInterface,
 				return
 			}
 
-			if resp.State != lnrpc.WalletState_WAITING_TO_START {
+			if predicate(resp.State) {
 				close(started)
 				return
 			}