peer: add test for createHtlcValidator

peer/brontide_test.go

@@ -6,6 +6,7 @@ import (
 	"testing"
 	"time"
 
+	"github.com/btcsuite/btcd/btcec/v2"
 	"github.com/btcsuite/btcd/btcutil"
 	"github.com/btcsuite/btcd/chaincfg"
 	"github.com/btcsuite/btcd/txscript"
@@ -19,6 +20,9 @@ import (
 	"github.com/lightningnetwork/lnd/lnwallet"
 	"github.com/lightningnetwork/lnd/lnwallet/chancloser"
 	"github.com/lightningnetwork/lnd/lnwire"
+	"github.com/lightningnetwork/lnd/routing/route"
+	"github.com/lightningnetwork/lnd/tlv"
+	"github.com/stretchr/testify/mock"
 	"github.com/stretchr/testify/require"
 )
 
@@ -1466,3 +1470,210 @@ func TestRemovePendingChannel(t *testing.T) {
 
 	require.NoError(t, err)
 }
+
+// mockAuxTrafficShaper is a mock implementation of htlcswitch.AuxTrafficShaper
+// for testing the createHtlcValidator function.
+type mockAuxTrafficShaper struct {
+	mock.Mock
+}
+
+// ShouldHandleTraffic returns the configured mock values.
+func (m *mockAuxTrafficShaper) ShouldHandleTraffic(
+	cid lnwire.ShortChannelID,
+	fundingBlob, htlcBlob fn.Option[tlv.Blob]) (bool, error) {
+
+	args := m.Called(cid, fundingBlob, htlcBlob)
+	return args.Bool(0), args.Error(1)
+}
+
+// PaymentBandwidth returns the configured mock values.
+func (m *mockAuxTrafficShaper) PaymentBandwidth(fundingBlob, htlcBlob,
+	commitmentBlob fn.Option[tlv.Blob], linkBandwidth,
+	htlcAmt lnwire.MilliSatoshi, htlcView lnwallet.AuxHtlcView,
+	peer route.Vertex) (lnwire.MilliSatoshi, error) {
+
+	args := m.Called(
+		fundingBlob, htlcBlob, commitmentBlob, linkBandwidth,
+		htlcAmt, htlcView, peer,
+	)
+
+	bw, _ := args.Get(0).(lnwire.MilliSatoshi)
+
+	return bw, args.Error(1)
+}
+
+// ProduceHtlcExtraData is part of the AuxTrafficShaper interface.
+func (m *mockAuxTrafficShaper) ProduceHtlcExtraData(
+	totalAmount lnwire.MilliSatoshi,
+	htlcCustomRecords lnwire.CustomRecords,
+	peer route.Vertex) (lnwire.MilliSatoshi, lnwire.CustomRecords,
+	error) {
+
+	args := m.Called(totalAmount, htlcCustomRecords, peer)
+
+	amt, _ := args.Get(0).(lnwire.MilliSatoshi)
+	records, _ := args.Get(1).(lnwire.CustomRecords)
+
+	return amt, records, args.Error(2)
+}
+
+// IsCustomHTLC is part of the AuxTrafficShaper interface.
+func (m *mockAuxTrafficShaper) IsCustomHTLC(
+	htlcRecords lnwire.CustomRecords) bool {
+
+	args := m.Called(htlcRecords)
+	return args.Bool(0)
+}
+
+// Compile-time check that mockAuxTrafficShaper implements AuxTrafficShaper.
+var _ htlcswitch.AuxTrafficShaper = (*mockAuxTrafficShaper)(nil)
+
+// TestCreateHtlcValidator tests that the HTLC validator created by
+// createHtlcValidator respects the ShouldHandleTraffic check. When
+// ShouldHandleTraffic returns false, the validator should return nil without
+// calling PaymentBandwidth.
+func TestCreateHtlcValidator(t *testing.T) {
+	t.Parallel()
+
+	// Create a minimal Brontide with just the identity key set.
+	privKey, err := btcec.NewPrivateKey()
+	require.NoError(t, err)
+
+	peer := &Brontide{
+		cfg: Config{
+			Addr: &lnwire.NetAddress{
+				IdentityKey: privKey.PubKey(),
+			},
+		},
+	}
+
+	// Create a mock channel with minimal required fields.
+	dbChan := &channeldb.OpenChannel{
+		ShortChannelID: lnwire.NewShortChanIDFromInt(123),
+	}
+
+	anyArg := mock.Anything
+
+	testCases := []struct {
+		name        string
+		setupMock   func(*mockAuxTrafficShaper)
+		htlcAmount  lnwire.MilliSatoshi
+		linkBw      lnwire.MilliSatoshi
+		expectError bool
+	}{
+		{
+			name: "non-custom channel skips check",
+			setupMock: func(m *mockAuxTrafficShaper) {
+				m.On(
+					"ShouldHandleTraffic",
+					anyArg, anyArg, anyArg,
+				).Return(false, nil)
+			},
+			htlcAmount:  1000,
+			linkBw:      5000,
+			expectError: false,
+		},
+		{
+			name: "sufficient bandwidth",
+			setupMock: func(m *mockAuxTrafficShaper) {
+				m.On(
+					"ShouldHandleTraffic",
+					anyArg, anyArg, anyArg,
+				).Return(true, nil)
+				m.On(
+					"PaymentBandwidth",
+					anyArg, anyArg, anyArg,
+					anyArg, anyArg, anyArg,
+					anyArg,
+				).Return(
+					lnwire.MilliSatoshi(10000),
+					nil,
+				)
+			},
+			htlcAmount:  1000,
+			linkBw:      5000,
+			expectError: false,
+		},
+		{
+			name: "insufficient bandwidth",
+			setupMock: func(m *mockAuxTrafficShaper) {
+				m.On(
+					"ShouldHandleTraffic",
+					anyArg, anyArg, anyArg,
+				).Return(true, nil)
+				m.On(
+					"PaymentBandwidth",
+					anyArg, anyArg, anyArg,
+					anyArg, anyArg, anyArg,
+					anyArg,
+				).Return(
+					lnwire.MilliSatoshi(500),
+					nil,
+				)
+			},
+			htlcAmount:  1000,
+			linkBw:      5000,
+			expectError: true,
+		},
+		{
+			name: "ShouldHandleTraffic error",
+			setupMock: func(m *mockAuxTrafficShaper) {
+				m.On(
+					"ShouldHandleTraffic",
+					anyArg, anyArg, anyArg,
+				).Return(
+					false,
+					fmt.Errorf("shaper error"),
+				)
+			},
+			htlcAmount:  1000,
+			linkBw:      5000,
+			expectError: true,
+		},
+		{
+			name: "PaymentBandwidth error",
+			setupMock: func(m *mockAuxTrafficShaper) {
+				m.On(
+					"ShouldHandleTraffic",
+					anyArg, anyArg, anyArg,
+				).Return(true, nil)
+				m.On(
+					"PaymentBandwidth",
+					anyArg, anyArg, anyArg,
+					anyArg, anyArg, anyArg,
+					anyArg,
+				).Return(
+					lnwire.MilliSatoshi(0),
+					fmt.Errorf("bandwidth error"),
+				)
+			},
+			htlcAmount:  1000,
+			linkBw:      5000,
+			expectError: true,
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			m := &mockAuxTrafficShaper{}
+			tc.setupMock(m)
+
+			validator := peer.createHtlcValidator(
+				dbChan, m,
+			)
+
+			err := validator.ValidateHtlc(
+				tc.htlcAmount, tc.linkBw,
+				nil, lnwallet.AuxHtlcView{},
+			)
+
+			if tc.expectError {
+				require.Error(t, err)
+			} else {
+				require.NoError(t, err)
+			}
+
+			m.AssertExpectations(t)
+		})
+	}
+}