From 608d11dcbc6b0cfe4441951aeaf1a5655c2bbb06 Mon Sep 17 00:00:00 2001
From: eugene <elzeigel@gmail.com>
Date: Thu, 12 Aug 2021 12:34:30 -0400
Subject: [PATCH] peer: add mockUpdateHandler for use in mockMessageSwitch

---
 peer/brontide_test.go | 10 +++++-----
 peer/test_utils.go    | 42 +++++++++++++++++++++++++++++++++++++++---
 2 files changed, 44 insertions(+), 8 deletions(-)

diff --git a/peer/brontide_test.go b/peer/brontide_test.go
index 2141a945f..588090878 100644
--- a/peer/brontide_test.go
+++ b/peer/brontide_test.go
@@ -40,7 +40,7 @@ func TestPeerChannelClosureAcceptFeeResponder(t *testing.T) {
 	broadcastTxChan := make(chan *wire.MsgTx)
 
 	alicePeer, bobChan, cleanUp, err := createTestPeer(
-		notifier, broadcastTxChan, noUpdate,
+		notifier, broadcastTxChan, noUpdate, nil,
 	)
 	if err != nil {
 		t.Fatalf("unable to create test channels: %v", err)
@@ -143,7 +143,7 @@ func TestPeerChannelClosureAcceptFeeInitiator(t *testing.T) {
 	broadcastTxChan := make(chan *wire.MsgTx)
 
 	alicePeer, bobChan, cleanUp, err := createTestPeer(
-		notifier, broadcastTxChan, noUpdate,
+		notifier, broadcastTxChan, noUpdate, nil,
 	)
 	if err != nil {
 		t.Fatalf("unable to create test channels: %v", err)
@@ -265,7 +265,7 @@ func TestPeerChannelClosureFeeNegotiationsResponder(t *testing.T) {
 	broadcastTxChan := make(chan *wire.MsgTx)
 
 	alicePeer, bobChan, cleanUp, err := createTestPeer(
-		notifier, broadcastTxChan, noUpdate,
+		notifier, broadcastTxChan, noUpdate, nil,
 	)
 	if err != nil {
 		t.Fatalf("unable to create test channels: %v", err)
@@ -459,7 +459,7 @@ func TestPeerChannelClosureFeeNegotiationsInitiator(t *testing.T) {
 	broadcastTxChan := make(chan *wire.MsgTx)
 
 	alicePeer, bobChan, cleanUp, err := createTestPeer(
-		notifier, broadcastTxChan, noUpdate,
+		notifier, broadcastTxChan, noUpdate, nil,
 	)
 	if err != nil {
 		t.Fatalf("unable to create test channels: %v", err)
@@ -795,7 +795,7 @@ func TestCustomShutdownScript(t *testing.T) {
 
 			// Open a channel.
 			alicePeer, bobChan, cleanUp, err := createTestPeer(
-				notifier, broadcastTxChan, test.update,
+				notifier, broadcastTxChan, test.update, nil,
 			)
 			if err != nil {
 				t.Fatalf("unable to create test channels: %v", err)
diff --git a/peer/test_utils.go b/peer/test_utils.go
index b6cf16b06..1d63a67bf 100644
--- a/peer/test_utils.go
+++ b/peer/test_utils.go
@@ -54,7 +54,8 @@ var noUpdate = func(a, b *channeldb.OpenChannel) {}
 // an updateChan function which can be used to modify the default values on
 // the channel states for each peer.
 func createTestPeer(notifier chainntnfs.ChainNotifier,
-	publTx chan *wire.MsgTx, updateChan func(a, b *channeldb.OpenChannel)) (
+	publTx chan *wire.MsgTx, updateChan func(a, b *channeldb.OpenChannel),
+	mockSwitch *mockMessageSwitch) (
 	*Brontide, *lnwallet.LightningChannel, func(), error) {
 
 	aliceKeyPriv, aliceKeyPub := btcec.PrivKeyFromBytes(
@@ -306,7 +307,11 @@ func createTestPeer(notifier chainntnfs.ChainNotifier,
 		},
 	}
 
-	htlcSwitch := &mockMessageSwitch{}
+	// If mockSwitch is not set by the caller, set it to the default as the
+	// caller does not need to control it.
+	if mockSwitch == nil {
+		mockSwitch = &mockMessageSwitch{}
+	}
 
 	nodeSignerAlice := netann.NewNodeSigner(aliceKeySigner)
 
@@ -349,7 +354,7 @@ func createTestPeer(notifier chainntnfs.ChainNotifier,
 		PubKeyBytes: pubKey,
 		ErrorBuffer: errBuffer,
 		ChainIO:     chainIO,
-		Switch:      htlcSwitch,
+		Switch:      mockSwitch,
 
 		ChanActiveTimeout: chanActiveTimeout,
 		InterceptSwitch:   htlcswitch.NewInterceptableSwitch(nil),
@@ -404,6 +409,37 @@ func (m *mockMessageSwitch) GetLinksByInterface(pub [33]byte) (
 	return nil, nil
 }
 
+// mockUpdateHandler is a mock implementation of the ChannelUpdateHandler
+// interface. It is used in mockMessageSwitch's GetLinksByInterface method.
+type mockUpdateHandler struct {
+	cid lnwire.ChannelID
+}
+
+// newMockUpdateHandler creates a new mockUpdateHandler.
+func newMockUpdateHandler(cid lnwire.ChannelID) *mockUpdateHandler {
+	return &mockUpdateHandler{
+		cid: cid,
+	}
+}
+
+// HandleChannelUpdate currently does nothing.
+func (m *mockUpdateHandler) HandleChannelUpdate(msg lnwire.Message) {}
+
+// ChanID returns the mockUpdateHandler's cid.
+func (m *mockUpdateHandler) ChanID() lnwire.ChannelID { return m.cid }
+
+// Bandwidth currently returns a dummy value.
+func (m *mockUpdateHandler) Bandwidth() lnwire.MilliSatoshi { return 0 }
+
+// EligibleToForward currently returns a dummy value.
+func (m *mockUpdateHandler) EligibleToForward() bool { return false }
+
+// MayAddOutgoingHtlc currently returns nil.
+func (m *mockUpdateHandler) MayAddOutgoingHtlc() error { return nil }
+
+// ShutdownIfChannelClean currently returns nil.
+func (m *mockUpdateHandler) ShutdownIfChannelClean() error { return nil }
+
 type mockMessageConn struct {
 	t *testing.T