sweep: introduce BudgetAggregator to cluster inputs by deadlines

This commit adds `BudgetAggregator` as a new implementation of
`UtxoAggregator`. This aggregator will group inputs by their deadline
heights and create input sets that can be used directly by the fee
bumper for fee calculations.

input/mocks.go

@@ -123,3 +123,48 @@ func (m *MockInput) UnconfParent() *TxInfo {
 
 	return info.(*TxInfo)
 }
+
+// MockWitnessType implements the `WitnessType` interface and is used by other
+// packages for mock testing.
+type MockWitnessType struct {
+	mock.Mock
+}
+
+// Compile time assertion that MockWitnessType implements WitnessType.
+var _ WitnessType = (*MockWitnessType)(nil)
+
+// String returns a human readable version of the WitnessType.
+func (m *MockWitnessType) String() string {
+	args := m.Called()
+
+	return args.String(0)
+}
+
+// WitnessGenerator will return a WitnessGenerator function that an output uses
+// to generate the witness and optionally the sigScript for a sweep
+// transaction.
+func (m *MockWitnessType) WitnessGenerator(signer Signer,
+	descriptor *SignDescriptor) WitnessGenerator {
+
+	args := m.Called()
+
+	return args.Get(0).(WitnessGenerator)
+}
+
+// SizeUpperBound returns the maximum length of the witness of this WitnessType
+// if it would be included in a tx. It also returns if the output itself is a
+// nested p2sh output, if so then we need to take into account the extra
+// sigScript data size.
+func (m *MockWitnessType) SizeUpperBound() (int, bool, error) {
+	args := m.Called()
+
+	return args.Int(0), args.Bool(1), args.Error(2)
+}
+
+// AddWeightEstimation adds the estimated size of the witness in bytes to the
+// given weight estimator.
+func (m *MockWitnessType) AddWeightEstimation(e *TxWeightEstimator) error {
+	args := m.Called()
+
+	return args.Error(0)
+}

sweep/aggregator.go

@@ -3,7 +3,10 @@ package sweep
 import (
 	"sort"
 
+	"github.com/btcsuite/btcd/btcutil"
 	"github.com/btcsuite/btcd/wire"
+	"github.com/lightningnetwork/lnd/fn"
+	"github.com/lightningnetwork/lnd/lnwallet"
 	"github.com/lightningnetwork/lnd/lnwallet/chainfee"
 )
 
@@ -461,3 +464,232 @@ func zipClusters(as, bs []inputCluster) []inputCluster {
 
 	return finalClusters
 }
+
+// BudgetAggregator is a budget-based aggregator that creates clusters based on
+// deadlines and budgets of inputs.
+type BudgetAggregator struct {
+	// estimator is used when crafting sweep transactions to estimate the
+	// necessary fee relative to the expected size of the sweep
+	// transaction.
+	estimator chainfee.Estimator
+
+	// maxInputs specifies the maximum number of inputs allowed in a single
+	// sweep tx.
+	maxInputs uint32
+}
+
+// Compile-time constraint to ensure BudgetAggregator implements UtxoAggregator.
+var _ UtxoAggregator = (*BudgetAggregator)(nil)
+
+// NewBudgetAggregator creates a new instance of a BudgetAggregator.
+func NewBudgetAggregator(estimator chainfee.Estimator,
+	maxInputs uint32) *BudgetAggregator {
+
+	return &BudgetAggregator{
+		estimator: estimator,
+		maxInputs: maxInputs,
+	}
+}
+
+// clusterGroup defines an alias for a set of inputs that are to be grouped.
+type clusterGroup map[fn.Option[int32]][]pendingInput
+
+// ClusterInputs creates a list of input sets from pending inputs.
+// 1. filter out inputs whose budget cannot cover min relay fee.
+// 2. group the inputs into clusters based on their deadline height.
+// 3. sort the inputs in each cluster by their budget.
+// 4. optionally split a cluster if it exceeds the max input limit.
+// 5. create input sets from each of the clusters.
+func (b *BudgetAggregator) ClusterInputs(inputs pendingInputs) []InputSet {
+	// Filter out inputs that have a budget below min relay fee.
+	filteredInputs := b.filterInputs(inputs)
+
+	// Create clusters to group inputs based on their deadline height.
+	clusters := make(clusterGroup, len(filteredInputs))
+
+	// Iterate all the inputs and group them based on their specified
+	// deadline heights.
+	for _, input := range filteredInputs {
+		height := input.params.DeadlineHeight
+		cluster, ok := clusters[height]
+		if !ok {
+			cluster = make([]pendingInput, 0)
+		}
+
+		cluster = append(cluster, *input)
+		clusters[height] = cluster
+	}
+
+	// Now that we have the clusters, we can create the input sets.
+	//
+	// NOTE: cannot pre-allocate the slice since we don't know the number
+	// of input sets in advance.
+	inputSets := make([]InputSet, 0)
+	for _, cluster := range clusters {
+		// Sort the inputs by their economical value.
+		sortedInputs := b.sortInputs(cluster)
+
+		// Create input sets from the cluster.
+		sets := b.createInputSets(sortedInputs)
+		inputSets = append(inputSets, sets...)
+	}
+
+	return inputSets
+}
+
+// createInputSet takes a set of inputs which share the same deadline height
+// and turns them into a list of `InputSet`, each set is then used to create a
+// sweep transaction.
+func (b *BudgetAggregator) createInputSets(inputs []pendingInput) []InputSet {
+	// sets holds the InputSets that we will return.
+	sets := make([]InputSet, 0)
+
+	// Copy the inputs to a new slice so we can modify it.
+	remainingInputs := make([]pendingInput, len(inputs))
+	copy(remainingInputs, inputs)
+
+	// If the number of inputs is greater than the max inputs allowed, we
+	// will split them into smaller clusters.
+	for uint32(len(remainingInputs)) > b.maxInputs {
+		log.Tracef("Cluster has %v inputs, max is %v, dividing...",
+			len(inputs), b.maxInputs)
+
+		// Copy the inputs to be put into the new set, and update the
+		// remaining inputs by removing currentInputs.
+		currentInputs := make([]pendingInput, b.maxInputs)
+		copy(currentInputs, remainingInputs[:b.maxInputs])
+		remainingInputs = remainingInputs[b.maxInputs:]
+
+		// Create an InputSet using the max allowed number of inputs.
+		set, err := NewBudgetInputSet(currentInputs)
+		if err != nil {
+			log.Errorf("unable to create input set: %v", err)
+
+			continue
+		}
+
+		sets = append(sets, set)
+	}
+
+	// Create an InputSet from the remaining inputs.
+	if len(remainingInputs) > 0 {
+		set, err := NewBudgetInputSet(remainingInputs)
+		if err != nil {
+			log.Errorf("unable to create input set: %v", err)
+			return nil
+		}
+
+		sets = append(sets, set)
+	}
+
+	return sets
+}
+
+// filterInputs filters out inputs that have a budget below the min relay fee
+// or have a required output that's below the dust.
+func (b *BudgetAggregator) filterInputs(inputs pendingInputs) pendingInputs {
+	// Get the current min relay fee for this round.
+	minFeeRate := b.estimator.RelayFeePerKW()
+
+	// filterInputs stores a map of inputs that has a budget that at least
+	// can pay the minimal fee.
+	filteredInputs := make(pendingInputs, len(inputs))
+
+	// Iterate all the inputs and filter out the ones whose budget cannot
+	// cover the min fee.
+	for _, pi := range inputs {
+		op := pi.OutPoint()
+
+		// Get the size and skip if there's an error.
+		size, _, err := pi.WitnessType().SizeUpperBound()
+		if err != nil {
+			log.Warnf("Skipped input=%v: cannot get its size: %v",
+				op, err)
+
+			continue
+		}
+
+		// Skip inputs that has too little budget.
+		minFee := minFeeRate.FeeForWeight(int64(size))
+		if pi.params.Budget < minFee {
+			log.Warnf("Skipped input=%v: has budget=%v, but the "+
+				"min fee requires %v", op, pi.params.Budget,
+				minFee)
+
+			continue
+		}
+
+		// If the input comes with a required tx out that is below
+		// dust, we won't add it.
+		//
+		// NOTE: only HtlcSecondLevelAnchorInput returns non-nil
+		// RequiredTxOut.
+		reqOut := pi.RequiredTxOut()
+		if reqOut != nil {
+			if isDustOutput(reqOut) {
+				log.Errorf("Rejected input=%v due to dust "+
+					"required output=%v", op, reqOut.Value)
+
+				continue
+			}
+		}
+
+		filteredInputs[*op] = pi
+	}
+
+	return filteredInputs
+}
+
+// sortInputs sorts the inputs based on their economical value.
+//
+// NOTE: besides the forced inputs, the sorting won't make any difference
+// because all the inputs are added to the same set. The exception is when the
+// number of inputs exceeds the maxInputs limit, it requires us to split them
+// into smaller clusters. In that case, the sorting will make a difference as
+// the budgets of the clusters will be different.
+func (b *BudgetAggregator) sortInputs(inputs []pendingInput) []pendingInput {
+	// sortedInputs is the final list of inputs sorted by their economical
+	// value.
+	sortedInputs := make([]pendingInput, 0, len(inputs))
+
+	// Copy the inputs.
+	sortedInputs = append(sortedInputs, inputs...)
+
+	// Sort the inputs based on their budgets.
+	//
+	// NOTE: We can implement more sophisticated algorithm as the budget
+	// left is a function f(minFeeRate, size) = b1 - s1 * r > b2 - s2 * r,
+	// where b1 and b2 are budgets, s1 and s2 are sizes of the inputs.
+	sort.Slice(sortedInputs, func(i, j int) bool {
+		left := sortedInputs[i].params.Budget
+		right := sortedInputs[j].params.Budget
+
+		// Make sure forced inputs are always put in the front.
+		leftForce := sortedInputs[i].params.Force
+		rightForce := sortedInputs[j].params.Force
+
+		// If both are forced inputs, we return the one with the higher
+		// budget. If neither are forced inputs, we also return the one
+		// with the higher budget.
+		if leftForce == rightForce {
+			return left > right
+		}
+
+		// Otherwise, it's either the left or the right is forced. We
+		// can simply return `leftForce` here as, if it's true, the
+		// left is forced and should be put in the front. Otherwise,
+		// the right is forced and should be put in the front.
+		return leftForce
+	})
+
+	return sortedInputs
+}
+
+// isDustOutput checks if the given output is considered as dust.
+func isDustOutput(output *wire.TxOut) bool {
+	// Fetch the dust limit for this output.
+	dustLimit := lnwallet.DustLimitForSize(len(output.PkScript))
+
+	// If the output is below the dust limit, we consider it dust.
+	return btcutil.Amount(output.Value) < dustLimit
+}

sweep/aggregator_test.go

@@ -1,14 +1,17 @@
 package sweep
 
 import (
+	"bytes"
 	"errors"
 	"reflect"
 	"sort"
 	"testing"
 
+	"github.com/btcsuite/btcd/btcutil"
 	"github.com/btcsuite/btcd/chaincfg/chainhash"
 	"github.com/btcsuite/btcd/wire"
 	"github.com/davecgh/go-spew/spew"
+	"github.com/lightningnetwork/lnd/fn"
 	"github.com/lightningnetwork/lnd/input"
 	"github.com/lightningnetwork/lnd/lnwallet/chainfee"
 	"github.com/stretchr/testify/require"
@@ -421,3 +424,510 @@ func TestClusterByLockTime(t *testing.T) {
 		})
 	}
 }
+
+// TestBudgetAggregatorFilterInputs checks that inputs with low budget are
+// filtered out.
+func TestBudgetAggregatorFilterInputs(t *testing.T) {
+	t.Parallel()
+
+	// Create a mock fee estimator.
+	estimator := &chainfee.MockEstimator{}
+	defer estimator.AssertExpectations(t)
+
+	// Create a mock WitnessType that always return an error when trying to
+	// get its size upper bound.
+	wtErr := &input.MockWitnessType{}
+	defer wtErr.AssertExpectations(t)
+
+	// Mock the `SizeUpperBound` method to return an error exactly once.
+	dummyErr := errors.New("dummy error")
+	wtErr.On("SizeUpperBound").Return(0, false, dummyErr).Once()
+
+	// Create a mock WitnessType that gives the size.
+	wt := &input.MockWitnessType{}
+	defer wt.AssertExpectations(t)
+
+	// Mock the `SizeUpperBound` method to return the size four times.
+	const wtSize = 100
+	wt.On("SizeUpperBound").Return(wtSize, true, nil).Times(4)
+
+	// Create a mock input that will be filtered out due to error.
+	inpErr := &input.MockInput{}
+	defer inpErr.AssertExpectations(t)
+
+	// Mock the `WitnessType` method to return the erroring witness type.
+	inpErr.On("WitnessType").Return(wtErr).Once()
+
+	// Mock the `OutPoint` method to return a unique outpoint.
+	opErr := wire.OutPoint{Hash: chainhash.Hash{1}}
+	inpErr.On("OutPoint").Return(&opErr).Once()
+
+	// Mock the estimator to return a constant fee rate.
+	const minFeeRate = chainfee.SatPerKWeight(1000)
+	estimator.On("RelayFeePerKW").Return(minFeeRate).Once()
+
+	var (
+		// Define three budget values, one below the min fee rate, one
+		// above and one equal to it.
+		budgetLow   = minFeeRate.FeeForWeight(wtSize) - 1
+		budgetEqual = minFeeRate.FeeForWeight(wtSize)
+		budgetHigh  = minFeeRate.FeeForWeight(wtSize) + 1
+
+		// Define three outpoints with different budget values.
+		opLow   = wire.OutPoint{Hash: chainhash.Hash{2}}
+		opEqual = wire.OutPoint{Hash: chainhash.Hash{3}}
+		opHigh  = wire.OutPoint{Hash: chainhash.Hash{4}}
+
+		// Define an outpoint that has a dust required output.
+		opDust = wire.OutPoint{Hash: chainhash.Hash{5}}
+	)
+
+	// Create three mock inputs.
+	inpLow := &input.MockInput{}
+	defer inpLow.AssertExpectations(t)
+
+	inpEqual := &input.MockInput{}
+	defer inpEqual.AssertExpectations(t)
+
+	inpHigh := &input.MockInput{}
+	defer inpHigh.AssertExpectations(t)
+
+	inpDust := &input.MockInput{}
+	defer inpDust.AssertExpectations(t)
+
+	// Mock the `WitnessType` method to return the witness type.
+	inpLow.On("WitnessType").Return(wt)
+	inpEqual.On("WitnessType").Return(wt)
+	inpHigh.On("WitnessType").Return(wt)
+	inpDust.On("WitnessType").Return(wt)
+
+	// Mock the `OutPoint` method to return the unique outpoint.
+	inpLow.On("OutPoint").Return(&opLow)
+	inpEqual.On("OutPoint").Return(&opEqual)
+	inpHigh.On("OutPoint").Return(&opHigh)
+	inpDust.On("OutPoint").Return(&opDust)
+
+	// Mock the `RequiredTxOut` to return nil.
+	inpEqual.On("RequiredTxOut").Return(nil)
+	inpHigh.On("RequiredTxOut").Return(nil)
+
+	// Mock the dust required output.
+	inpDust.On("RequiredTxOut").Return(&wire.TxOut{
+		Value:    0,
+		PkScript: bytes.Repeat([]byte{0}, input.P2WSHSize),
+	})
+
+	// Create testing pending inputs.
+	inputs := pendingInputs{
+		// The first input will be filtered out due to the error.
+		opErr: &pendingInput{
+			Input: inpErr,
+		},
+
+		// The second input will be filtered out due to the budget.
+		opLow: &pendingInput{
+			Input:  inpLow,
+			params: Params{Budget: budgetLow},
+		},
+
+		// The third input will be included.
+		opEqual: &pendingInput{
+			Input:  inpEqual,
+			params: Params{Budget: budgetEqual},
+		},
+
+		// The fourth input will be included.
+		opHigh: &pendingInput{
+			Input:  inpHigh,
+			params: Params{Budget: budgetHigh},
+		},
+
+		// The fifth input will be filtered out due to the dust
+		// required.
+		opDust: &pendingInput{
+			Input:  inpDust,
+			params: Params{Budget: budgetHigh},
+		},
+	}
+
+	// Init the budget aggregator with the mocked estimator and zero max
+	// num of inputs.
+	b := NewBudgetAggregator(estimator, 0)
+
+	// Call the method under test.
+	result := b.filterInputs(inputs)
+
+	// Validate the expected inputs are returned.
+	require.Len(t, result, 2)
+
+	// We expect only the inputs with budget equal or above the min fee to
+	// be included.
+	require.Contains(t, result, opEqual)
+	require.Contains(t, result, opHigh)
+}
+
+// TestBudgetAggregatorSortInputs checks that inputs are sorted by based on
+// their budgets and force flag.
+func TestBudgetAggregatorSortInputs(t *testing.T) {
+	t.Parallel()
+
+	var (
+		// Create two budgets.
+		budgetLow   = btcutil.Amount(1000)
+		budgetHight = budgetLow + btcutil.Amount(1000)
+	)
+
+	// Create an input with the low budget but forced.
+	inputLowForce := pendingInput{
+		params: Params{
+			Budget: budgetLow,
+			Force:  true,
+		},
+	}
+
+	// Create an input with the low budget.
+	inputLow := pendingInput{
+		params: Params{
+			Budget: budgetLow,
+		},
+	}
+
+	// Create an input with the high budget and forced.
+	inputHighForce := pendingInput{
+		params: Params{
+			Budget: budgetHight,
+			Force:  true,
+		},
+	}
+
+	// Create an input with the high budget.
+	inputHigh := pendingInput{
+		params: Params{
+			Budget: budgetHight,
+		},
+	}
+
+	// Create a testing pending inputs.
+	inputs := []pendingInput{
+		inputLowForce,
+		inputLow,
+		inputHighForce,
+		inputHigh,
+	}
+
+	// Init the budget aggregator with zero max num of inputs.
+	b := NewBudgetAggregator(nil, 0)
+
+	// Call the method under test.
+	result := b.sortInputs(inputs)
+	require.Len(t, result, 4)
+
+	// The first input should be the forced input with the high budget.
+	require.Equal(t, inputHighForce, result[0])
+
+	// The second input should be the forced input with the low budget.
+	require.Equal(t, inputLowForce, result[1])
+
+	// The third input should be the input with the high budget.
+	require.Equal(t, inputHigh, result[2])
+
+	// The fourth input should be the input with the low budget.
+	require.Equal(t, inputLow, result[3])
+}
+
+// TestBudgetAggregatorCreateInputSets checks that the budget aggregator
+// creates input sets when the number of inputs exceeds the max number
+// configed.
+func TestBudgetAggregatorCreateInputSets(t *testing.T) {
+	t.Parallel()
+
+	// Create mocks input that doesn't have required outputs.
+	mockInput1 := &input.MockInput{}
+	defer mockInput1.AssertExpectations(t)
+	mockInput2 := &input.MockInput{}
+	defer mockInput2.AssertExpectations(t)
+	mockInput3 := &input.MockInput{}
+	defer mockInput3.AssertExpectations(t)
+	mockInput4 := &input.MockInput{}
+	defer mockInput4.AssertExpectations(t)
+
+	// Create testing pending inputs.
+	pi1 := pendingInput{
+		Input: mockInput1,
+		params: Params{
+			DeadlineHeight: fn.Some(int32(1)),
+		},
+	}
+	pi2 := pendingInput{
+		Input: mockInput2,
+		params: Params{
+			DeadlineHeight: fn.Some(int32(1)),
+		},
+	}
+	pi3 := pendingInput{
+		Input: mockInput3,
+		params: Params{
+			DeadlineHeight: fn.Some(int32(1)),
+		},
+	}
+	pi4 := pendingInput{
+		Input: mockInput4,
+		params: Params{
+			// This input has a deadline height that is different
+			// from the other inputs. When grouped with other
+			// inputs, it will cause an error to be returned.
+			DeadlineHeight: fn.Some(int32(2)),
+		},
+	}
+
+	// Create a budget aggregator with max number of inputs set to 2.
+	b := NewBudgetAggregator(nil, 2)
+
+	// Create test cases.
+	testCases := []struct {
+		name            string
+		inputs          []pendingInput
+		setupMock       func()
+		expectedNumSets int
+	}{
+		{
+			// When the number of inputs is below the max, a single
+			// input set is returned.
+			name:   "num inputs below max",
+			inputs: []pendingInput{pi1},
+			setupMock: func() {
+				// Mock methods used in loggings.
+				mockInput1.On("WitnessType").Return(
+					input.CommitmentAnchor)
+				mockInput1.On("OutPoint").Return(
+					&wire.OutPoint{Hash: chainhash.Hash{1}})
+			},
+			expectedNumSets: 1,
+		},
+		{
+			// When the number of inputs is equal to the max, a
+			// single input set is returned.
+			name:   "num inputs equal to max",
+			inputs: []pendingInput{pi1, pi2},
+			setupMock: func() {
+				// Mock methods used in loggings.
+				mockInput1.On("WitnessType").Return(
+					input.CommitmentAnchor)
+				mockInput2.On("WitnessType").Return(
+					input.CommitmentAnchor)
+
+				mockInput1.On("OutPoint").Return(
+					&wire.OutPoint{Hash: chainhash.Hash{1}})
+				mockInput2.On("OutPoint").Return(
+					&wire.OutPoint{Hash: chainhash.Hash{2}})
+			},
+			expectedNumSets: 1,
+		},
+		{
+			// When the number of inputs is above the max, multiple
+			// input sets are returned.
+			name:   "num inputs above max",
+			inputs: []pendingInput{pi1, pi2, pi3},
+			setupMock: func() {
+				// Mock methods used in loggings.
+				mockInput1.On("WitnessType").Return(
+					input.CommitmentAnchor)
+				mockInput2.On("WitnessType").Return(
+					input.CommitmentAnchor)
+				mockInput3.On("WitnessType").Return(
+					input.CommitmentAnchor)
+
+				mockInput1.On("OutPoint").Return(
+					&wire.OutPoint{Hash: chainhash.Hash{1}})
+				mockInput2.On("OutPoint").Return(
+					&wire.OutPoint{Hash: chainhash.Hash{2}})
+				mockInput3.On("OutPoint").Return(
+					&wire.OutPoint{Hash: chainhash.Hash{3}})
+			},
+			expectedNumSets: 2,
+		},
+		{
+			// When the number of inputs is above the max, but an
+			// error is returned from creating the first set, it
+			// shouldn't affect the remaining inputs.
+			name:   "num inputs above max with error",
+			inputs: []pendingInput{pi1, pi4, pi3},
+			setupMock: func() {
+				// Mock methods used in loggings.
+				mockInput1.On("WitnessType").Return(
+					input.CommitmentAnchor)
+				mockInput3.On("WitnessType").Return(
+					input.CommitmentAnchor)
+
+				mockInput1.On("OutPoint").Return(
+					&wire.OutPoint{Hash: chainhash.Hash{1}})
+				mockInput3.On("OutPoint").Return(
+					&wire.OutPoint{Hash: chainhash.Hash{3}})
+				mockInput4.On("OutPoint").Return(
+					&wire.OutPoint{Hash: chainhash.Hash{2}})
+			},
+			expectedNumSets: 1,
+		},
+	}
+
+	// Iterate over the test cases.
+	for _, tc := range testCases {
+		tc := tc
+
+		t.Run(tc.name, func(t *testing.T) {
+			// Setup the mocks.
+			tc.setupMock()
+
+			// Call the method under test.
+			result := b.createInputSets(tc.inputs)
+
+			// Validate the expected number of input sets are
+			// returned.
+			require.Len(t, result, tc.expectedNumSets)
+		})
+	}
+}
+
+// TestBudgetInputSetClusterInputs checks that the budget aggregator clusters
+// inputs into input sets based on their deadline heights.
+func TestBudgetInputSetClusterInputs(t *testing.T) {
+	t.Parallel()
+
+	// Create a mock fee estimator.
+	estimator := &chainfee.MockEstimator{}
+	defer estimator.AssertExpectations(t)
+
+	// Create a mock WitnessType that gives the size.
+	wt := &input.MockWitnessType{}
+	defer wt.AssertExpectations(t)
+
+	// Mock the `SizeUpperBound` method to return the size six times since
+	// we are using nine inputs.
+	const wtSize = 100
+	wt.On("SizeUpperBound").Return(wtSize, true, nil).Times(9)
+	wt.On("String").Return("mock witness type")
+
+	// Mock the estimator to return a constant fee rate.
+	const minFeeRate = chainfee.SatPerKWeight(1000)
+	estimator.On("RelayFeePerKW").Return(minFeeRate).Once()
+
+	var (
+		// Define two budget values, one below the min fee rate and one
+		// above it.
+		budgetLow  = minFeeRate.FeeForWeight(wtSize) - 1
+		budgetHigh = minFeeRate.FeeForWeight(wtSize) + 1
+
+		// Create three deadline heights, which means there are three
+		// groups of inputs to be expected.
+		deadlineNone = fn.None[int32]()
+		deadline1    = fn.Some(int32(1))
+		deadline2    = fn.Some(int32(2))
+	)
+
+	// Create testing pending inputs.
+	inputs := make(pendingInputs)
+
+	// For each deadline height, create two inputs with different budgets,
+	// one below the min fee rate and one above it. We should see the lower
+	// one being filtered out.
+	for i, deadline := range []fn.Option[int32]{
+		deadlineNone, deadline1, deadline2,
+	} {
+		// Define three outpoints.
+		opLow := wire.OutPoint{
+			Hash:  chainhash.Hash{byte(i)},
+			Index: uint32(i),
+		}
+		opHigh1 := wire.OutPoint{
+			Hash:  chainhash.Hash{byte(i + 1000)},
+			Index: uint32(i + 1000),
+		}
+		opHigh2 := wire.OutPoint{
+			Hash:  chainhash.Hash{byte(i + 2000)},
+			Index: uint32(i + 2000),
+		}
+
+		// Create mock inputs.
+		inpLow := &input.MockInput{}
+		defer inpLow.AssertExpectations(t)
+
+		inpHigh1 := &input.MockInput{}
+		defer inpHigh1.AssertExpectations(t)
+
+		inpHigh2 := &input.MockInput{}
+		defer inpHigh2.AssertExpectations(t)
+
+		// Mock the `OutPoint` method to return the unique outpoint.
+		//
+		// We expect the low budget input to call this method once in
+		// `filterInputs`.
+		inpLow.On("OutPoint").Return(&opLow).Once()
+
+		// We expect the high budget input to call this method three
+		// times, one in `filterInputs` and one in `createInputSet`,
+		// and one in `NewBudgetInputSet`.
+		inpHigh1.On("OutPoint").Return(&opHigh1).Times(3)
+		inpHigh2.On("OutPoint").Return(&opHigh2).Times(3)
+
+		// Mock the `WitnessType` method to return the witness type.
+		inpLow.On("WitnessType").Return(wt)
+		inpHigh1.On("WitnessType").Return(wt)
+		inpHigh2.On("WitnessType").Return(wt)
+
+		// Mock the `RequiredTxOut` to return nil.
+		inpHigh1.On("RequiredTxOut").Return(nil)
+		inpHigh2.On("RequiredTxOut").Return(nil)
+
+		// Add the low input, which should be filtered out.
+		inputs[opLow] = &pendingInput{
+			Input: inpLow,
+			params: Params{
+				Budget:         budgetLow,
+				DeadlineHeight: deadline,
+			},
+		}
+
+		// Add the high inputs, which should be included.
+		inputs[opHigh1] = &pendingInput{
+			Input: inpHigh1,
+			params: Params{
+				Budget:         budgetHigh,
+				DeadlineHeight: deadline,
+			},
+		}
+		inputs[opHigh2] = &pendingInput{
+			Input: inpHigh2,
+			params: Params{
+				Budget:         budgetHigh,
+				DeadlineHeight: deadline,
+			},
+		}
+	}
+
+	// Create a budget aggregator with a max number of inputs set to 100.
+	b := NewBudgetAggregator(estimator, DefaultMaxInputsPerTx)
+
+	// Call the method under test.
+	result := b.ClusterInputs(inputs)
+
+	// We expect three input sets to be returned, one for each deadline.
+	require.Len(t, result, 3)
+
+	// Check each input set has exactly two inputs.
+	deadlines := make(map[fn.Option[int32]]struct{})
+	for _, set := range result {
+		// We expect two inputs in each set.
+		require.Len(t, set.Inputs(), 2)
+
+		// We expect each set to have the expected budget.
+		require.Equal(t, budgetHigh*2, set.Budget())
+
+		// Save the deadlines.
+		deadlines[set.DeadlineHeight()] = struct{}{}
+	}
+
+	// We expect to see all three deadlines.
+	require.Contains(t, deadlines, deadlineNone)
+	require.Contains(t, deadlines, deadline1)
+	require.Contains(t, deadlines, deadline2)
+}