From 0c6a1558d5bbab51c554764f0f642207578a3015 Mon Sep 17 00:00:00 2001
From: ffranr <robyn@lightning.engineering>
Date: Tue, 30 Apr 2024 17:51:43 +0100
Subject: [PATCH] lntest: add `HtlcModifier` support to node RPC harness

This commit enhances the itest LND node harness to include support for
the new `HtlcModifier` RPC endpoint.
At the same time we move another method to the correct file.
---
 lntest/harness.go      | 39 +++++++++++++++++++++++++++++++++++++--
 lntest/rpc/invoices.go | 24 ++++++++++++------------
 lntest/rpc/router.go   | 16 ++++++++++++++++
 3 files changed, 65 insertions(+), 14 deletions(-)

diff --git a/lntest/harness.go b/lntest/harness.go
index e8996e9a6..79c1f61ff 100644
--- a/lntest/harness.go
+++ b/lntest/harness.go
@@ -16,6 +16,7 @@ import (
 	"github.com/lightningnetwork/lnd/fn"
 	"github.com/lightningnetwork/lnd/kvdb/etcd"
 	"github.com/lightningnetwork/lnd/lnrpc"
+	"github.com/lightningnetwork/lnd/lnrpc/invoicesrpc"
 	"github.com/lightningnetwork/lnd/lnrpc/routerrpc"
 	"github.com/lightningnetwork/lnd/lnrpc/walletrpc"
 	"github.com/lightningnetwork/lnd/lntest/miner"
@@ -2071,8 +2072,42 @@ func (h *HarnessTest) ReceiveHtlcInterceptor(
 		require.Fail(h, "timeout", "timeout intercepting htlc")
 
 	case err := <-errChan:
-		require.Failf(h, "err from stream",
-			"received err from stream: %v", err)
+		require.Failf(h, "err from HTLC interceptor stream",
+			"received err from HTLC interceptor stream: %v", err)
+
+	case updateMsg := <-chanMsg:
+		return updateMsg
+	}
+
+	return nil
+}
+
+// ReceiveInvoiceHtlcModification waits until a message is received on the
+// invoice HTLC modifier stream or the timeout is reached.
+func (h *HarnessTest) ReceiveInvoiceHtlcModification(
+	stream rpc.InvoiceHtlcModifierClient) *invoicesrpc.HtlcModifyRequest {
+
+	chanMsg := make(chan *invoicesrpc.HtlcModifyRequest)
+	errChan := make(chan error)
+	go func() {
+		// Consume one message. This will block until the message is
+		// received.
+		resp, err := stream.Recv()
+		if err != nil {
+			errChan <- err
+			return
+		}
+		chanMsg <- resp
+	}()
+
+	select {
+	case <-time.After(DefaultTimeout):
+		require.Fail(h, "timeout", "timeout invoice HTLC modifier")
+
+	case err := <-errChan:
+		require.Failf(h, "err from invoice HTLC modifier stream",
+			"received err from invoice HTLC modifier stream: %v",
+			err)
 
 	case updateMsg := <-chanMsg:
 		return updateMsg
diff --git a/lntest/rpc/invoices.go b/lntest/rpc/invoices.go
index 5b771968e..a5c560897 100644
--- a/lntest/rpc/invoices.go
+++ b/lntest/rpc/invoices.go
@@ -5,7 +5,6 @@ import (
 
 	"github.com/lightningnetwork/lnd/lnrpc"
 	"github.com/lightningnetwork/lnd/lnrpc/invoicesrpc"
-	"github.com/lightningnetwork/lnd/lnrpc/routerrpc"
 )
 
 // =====================
@@ -84,18 +83,19 @@ func (h *HarnessRPC) SubscribeSingleInvoice(rHash []byte) SingleInvoiceClient {
 	return client
 }
 
-type TrackPaymentClient routerrpc.Router_TrackPaymentV2Client
+type InvoiceHtlcModifierClient invoicesrpc.Invoices_HtlcModifierClient
 
-// TrackPaymentV2 creates a subscription client for given invoice and
-// asserts its creation.
-func (h *HarnessRPC) TrackPaymentV2(payHash []byte) TrackPaymentClient {
-	req := &routerrpc.TrackPaymentRequest{PaymentHash: payHash}
+// InvoiceHtlcModifier makes an RPC call to the node's RouterClient and asserts.
+func (h *HarnessRPC) InvoiceHtlcModifier() (InvoiceHtlcModifierClient,
+	context.CancelFunc) {
 
-	// TrackPaymentV2 needs to have the context alive for the entire test
-	// case as the returned client will be used for send and receive events
-	// stream. Thus we use runCtx here instead of a timeout context.
-	client, err := h.Router.TrackPaymentV2(h.runCtx, req)
-	h.NoError(err, "TrackPaymentV2")
+	// InvoiceHtlcModifier needs to have the context alive for the entire
+	// test case as the returned client will be used for send and receive
+	// events stream. Therefore, we use cancel context here instead of a
+	// timeout context.
+	ctxt, cancel := context.WithCancel(h.runCtx)
+	resp, err := h.Invoice.HtlcModifier(ctxt)
+	h.NoError(err, "InvoiceHtlcModifier")
 
-	return client
+	return resp, cancel
 }
diff --git a/lntest/rpc/router.go b/lntest/rpc/router.go
index 7c3f2e7c4..707c43a11 100644
--- a/lntest/rpc/router.go
+++ b/lntest/rpc/router.go
@@ -267,3 +267,19 @@ func (h *HarnessRPC) TrackPayments(
 
 	return resp
 }
+
+type TrackPaymentClient routerrpc.Router_TrackPaymentV2Client
+
+// TrackPaymentV2 creates a subscription client for given invoice and
+// asserts its creation.
+func (h *HarnessRPC) TrackPaymentV2(payHash []byte) TrackPaymentClient {
+	req := &routerrpc.TrackPaymentRequest{PaymentHash: payHash}
+
+	// TrackPaymentV2 needs to have the context alive for the entire test
+	// case as the returned client will be used for send and receive events
+	// stream. Thus we use runCtx here instead of a timeout context.
+	client, err := h.Router.TrackPaymentV2(h.runCtx, req)
+	h.NoError(err, "TrackPaymentV2")
+
+	return client
+}