From 7e0c04fea566144164d5d85ffa7dc86f4faf48e7 Mon Sep 17 00:00:00 2001
From: bitromortac <bitromortac@protonmail.com>
Date: Mon, 2 Feb 2026 10:17:34 +0100
Subject: [PATCH] rules: add test to demo failure for closed channels

This also introduces the use of testify mock to check which calls were
made.
---
 rules/channel_restrictions_test.go | 75 ++++++++++++++++++++++--------
 rules/peer_restrictions_test.go    | 38 +++++++++------
 2 files changed, 78 insertions(+), 35 deletions(-)

diff --git a/rules/channel_restrictions_test.go b/rules/channel_restrictions_test.go
index 6102851c..fb14cc0d 100644
--- a/rules/channel_restrictions_test.go
+++ b/rules/channel_restrictions_test.go
@@ -12,6 +12,7 @@ import (
 	"github.com/lightninglabs/lightning-terminal/session"
 	"github.com/lightninglabs/lndclient"
 	"github.com/lightningnetwork/lnd/lnrpc"
+	"github.com/stretchr/testify/mock"
 	"github.com/stretchr/testify/require"
 )
 
@@ -37,22 +38,25 @@ func TestChannelRestrictCheckRequest(t *testing.T) {
 
 	ctx := context.Background()
 	mgr := NewChannelRestrictMgr()
-	cfg := &mockLndClient{
-		channels: []lndclient.ChannelInfo{
-			{
-				ChannelID:    chanID1,
-				ChannelPoint: chanPointStr1,
-			},
-			{
-				ChannelID:    chanID2,
-				ChannelPoint: chanPointStr2,
-			},
-			{
-				ChannelID:    chanID3,
-				ChannelPoint: chanPointStr3,
-			},
+	cfg := &mockLndClient{}
+	cfg.On(
+		"ListChannels", mock.Anything, mock.Anything, mock.Anything,
+		mock.Anything,
+	).Return([]lndclient.ChannelInfo{
+		{
+			ChannelID:    chanID1,
+			ChannelPoint: chanPointStr1,
 		},
-	}
+		{
+			ChannelID:    chanID2,
+			ChannelPoint: chanPointStr2,
+		},
+		{
+			ChannelID:    chanID3,
+			ChannelPoint: chanPointStr3,
+		},
+	}, nil)
+
 	enf, err := mgr.NewEnforcer(ctx, cfg, &ChannelRestrict{
 		DenyList: []uint64{
 			chanID1, chanID2,
@@ -149,18 +153,23 @@ func newTXID() ([]byte, uint32, error) {
 type mockLndClient struct {
 	lndclient.LightningClient
 	Config
-
-	channels []lndclient.ChannelInfo
+	mock.Mock
 }
 
 func (m *mockLndClient) GetLndClient() lndclient.LightningClient {
 	return m
 }
 
-func (m *mockLndClient) ListChannels(_ context.Context, _, _ bool,
-	_ ...lndclient.ListChannelsOption) ([]lndclient.ChannelInfo, error) {
+func (m *mockLndClient) ListChannels(ctx context.Context, public, active bool,
+	opts ...lndclient.ListChannelsOption) ([]lndclient.ChannelInfo, error) {
 
-	return m.channels, nil
+	args := m.Called(ctx, public, active, opts)
+
+	if args.Get(0) != nil {
+		return args.Get(0).([]lndclient.ChannelInfo), args.Error(1)
+	}
+
+	return nil, args.Error(1)
 }
 
 // TestChannelRestrictRealToPseudo tests that the ChannelRestrict's RealToPseudo
@@ -292,3 +301,29 @@ func TestChannelRestrictRealToPseudo(t *testing.T) {
 		})
 	}
 }
+
+// TestChannelRestrictResilience ensures that the ChannelRestrictEnforcer is
+// resilient to missing channels during initialization.
+func TestChannelRestrictResilience(t *testing.T) {
+	var (
+		ctx = context.Background()
+		mgr = NewChannelRestrictMgr()
+	)
+
+	chanID1, _ := firewalldb.NewPseudoUint64()
+
+	// Initially, LND has no channels.
+	cfg := &mockLndClient{}
+	cfg.On(
+		"ListChannels", mock.Anything, mock.Anything, mock.Anything,
+		mock.Anything,
+	).Return([]lndclient.ChannelInfo{}, nil)
+
+	// We create an enforcer that denies chanID1 (maybe a closed channel or
+	// generally unknown). This will be fixed in a future commit.
+	_, err := mgr.NewEnforcer(ctx, cfg, &ChannelRestrict{
+		DenyList: []uint64{chanID1},
+	})
+	require.ErrorContains(t, err, "invalid channel ID")
+	cfg.AssertExpectations(t)
+}
diff --git a/rules/peer_restrictions_test.go b/rules/peer_restrictions_test.go
index abfa3054..5b8fedb7 100644
--- a/rules/peer_restrictions_test.go
+++ b/rules/peer_restrictions_test.go
@@ -12,10 +12,12 @@ import (
 	"github.com/lightninglabs/lndclient"
 	"github.com/lightningnetwork/lnd/lnrpc"
 	"github.com/lightningnetwork/lnd/routing/route"
+	"github.com/stretchr/testify/mock"
 	"github.com/stretchr/testify/require"
 )
 
 // TestPeerRestrictCheckRequest ensures that the PeerRestrictEnforcer correctly
+
 // accepts or denys a request.
 func TestPeerRestrictCheckRequest(t *testing.T) {
 	txid1, index1, err := newTXID()
@@ -51,28 +53,31 @@ func TestPeerRestrictCheckRequest(t *testing.T) {
 
 	ctx := context.Background()
 	mgr := NewPeerRestrictMgr()
-	cfg := &mockLndClient{
-		channels: []lndclient.ChannelInfo{
-			{
-				ChannelPoint: chanPointStr1,
-				PubKeyBytes:  peerKey1,
-			},
-			{
-				ChannelPoint: chanPointStr2,
-				PubKeyBytes:  peerKey2,
-			},
-			{
-				ChannelPoint: chanPointStr3,
-				PubKeyBytes:  peerKey3,
-			},
+	cfg := &mockLndClient{}
+	cfg.On(
+		"ListChannels", mock.Anything, mock.Anything, mock.Anything,
+		mock.Anything,
+	).Return([]lndclient.ChannelInfo{
+		{
+			ChannelPoint: chanPointStr1,
+			PubKeyBytes:  peerKey1,
 		},
-	}
+		{
+			ChannelPoint: chanPointStr2,
+			PubKeyBytes:  peerKey2,
+		},
+		{
+			ChannelPoint: chanPointStr3,
+			PubKeyBytes:  peerKey3,
+		},
+	}, nil)
 
 	enf, err := mgr.NewEnforcer(ctx, cfg, &PeerRestrict{
 		DenyList: []string{
 			peerID1, peerID2,
 		},
 	})
+
 	require.NoError(t, err)
 
 	// A request for an irrelevant URI should be allowed.
@@ -198,6 +203,9 @@ func TestPeerRestrictCheckRequest(t *testing.T) {
 		ctx, "/lnrpc.Lightning/BatchOpenChannel", batchOpenReq,
 	)
 	require.NoError(t, err)
+
+	// Assert expected calls were made.
+	cfg.AssertExpectations(t)
 }
 
 // TestPeerRestrictionRealToPseudo tests that the PeerRestriction's RealToPseudo