From a088501e47c80b5f5506efc76e8bf11c6550ab09 Mon Sep 17 00:00:00 2001 From: yyforyongyu Date: Tue, 27 Feb 2024 21:54:48 +0800 Subject: [PATCH] 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 | 45 ++++ sweep/aggregator.go | 232 ++++++++++++++++++ sweep/aggregator_test.go | 510 +++++++++++++++++++++++++++++++++++++++ 3 files changed, 787 insertions(+) diff --git a/input/mocks.go b/input/mocks.go index 965489eff..1fe6eb765 100644 --- a/input/mocks.go +++ b/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) +} diff --git a/sweep/aggregator.go b/sweep/aggregator.go index 379ff9829..58ac51132 100644 --- a/sweep/aggregator.go +++ b/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 +} diff --git a/sweep/aggregator_test.go b/sweep/aggregator_test.go index 2058464ad..bee1db529 100644 --- a/sweep/aggregator_test.go +++ b/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) +}