From 3c4fd1b484735f40b4c0d040c5b8f2f6409b8259 Mon Sep 17 00:00:00 2001
From: yyforyongyu <yong2452@gmail.com>
Date: Wed, 19 Mar 2025 01:41:35 +0800
Subject: [PATCH] sweep: refactor `AddWalletInputs` by adding `addWalletInput`

A minor refactor to prepare for upcoming changes.
---
 sweep/tx_input_set.go      | 40 ++++++++++++++++++++++++++-----------
 sweep/tx_input_set_test.go | 41 ++++++++++++++++++++++++++++++++++++++
 2 files changed, 69 insertions(+), 12 deletions(-)

diff --git a/sweep/tx_input_set.go b/sweep/tx_input_set.go
index b80d52b0e..ff3b7f309 100644
--- a/sweep/tx_input_set.go
+++ b/sweep/tx_input_set.go
@@ -244,6 +244,28 @@ func (b *BudgetInputSet) addInput(input SweeperInput) {
 	b.inputs = append(b.inputs, &input)
 }
 
+// addWalletInput takes a wallet UTXO and adds it as an input to be used as
+// budget for the input set.
+func (b *BudgetInputSet) addWalletInput(utxo *lnwallet.Utxo) error {
+	input, err := createWalletTxInput(utxo)
+	if err != nil {
+		return err
+	}
+
+	pi := SweeperInput{
+		Input: input,
+		params: Params{
+			DeadlineHeight: fn.Some(b.deadlineHeight),
+		},
+	}
+	b.addInput(pi)
+
+	log.Debugf("Added wallet input to input set: op=%v, amt=%v",
+		pi.OutPoint(), utxo.Value)
+
+	return nil
+}
+
 // NeedWalletInput returns true if the input set needs more wallet inputs.
 //
 // A set may need wallet inputs when it has a required output or its total
@@ -326,6 +348,11 @@ func (b *BudgetInputSet) AddWalletInputs(wallet Wallet) error {
 		return fmt.Errorf("list unspent witness: %w", err)
 	}
 
+	// Exit early if there are no wallet UTXOs.
+	if len(utxos) == 0 {
+		return ErrNotEnoughInputs
+	}
+
 	// Sort the UTXOs by putting smaller values at the start of the slice
 	// to avoid locking large UTXO for sweeping.
 	//
@@ -342,22 +369,11 @@ func (b *BudgetInputSet) AddWalletInputs(wallet Wallet) error {
 
 	// Add wallet inputs to the set until the specified budget is covered.
 	for _, utxo := range utxos {
-		input, err := createWalletTxInput(utxo)
+		err := b.addWalletInput(utxo)
 		if err != nil {
 			return err
 		}
 
-		pi := SweeperInput{
-			Input: input,
-			params: Params{
-				DeadlineHeight: fn.Some(b.deadlineHeight),
-			},
-		}
-		b.addInput(pi)
-
-		log.Debugf("Added wallet input to input set: op=%v, amt=%v",
-			pi.OutPoint(), utxo.Value)
-
 		// Return if we've reached the minimum output amount.
 		if !b.NeedWalletInput() {
 			return nil
diff --git a/sweep/tx_input_set_test.go b/sweep/tx_input_set_test.go
index 73f056a96..0ea056e49 100644
--- a/sweep/tx_input_set_test.go
+++ b/sweep/tx_input_set_test.go
@@ -133,6 +133,47 @@ func TestBudgetInputSetAddInput(t *testing.T) {
 	require.Equal(t, btcutil.Amount(200), set.Budget())
 }
 
+// TestAddWalletInput asserts `addWalletInput` successfully converts a wallet
+// UTXO into a `SweeperInput` with the correct deadline.
+func TestAddWalletInput(t *testing.T) {
+	t.Parallel()
+
+	// Create a testing deadline.
+	deadline := int32(1000)
+
+	// Initialize an empty input set.
+	set := &BudgetInputSet{
+		deadlineHeight: deadline,
+	}
+
+	// Create an utxo with unknown address type to trigger an error.
+	utxo := &lnwallet.Utxo{
+		AddressType: lnwallet.UnknownAddressType,
+	}
+
+	// Check that the error is returned from addWalletInput.
+	err := set.addWalletInput(utxo)
+	require.Error(t, err)
+
+	// Create a wallet utxo.
+	utxo = &lnwallet.Utxo{
+		AddressType: lnwallet.WitnessPubKey,
+		Value:       1000,
+	}
+
+	// Check that no error is returned from addWalletInput.
+	err = set.addWalletInput(utxo)
+	require.NoError(t, err)
+
+	// Check the input has been added to the set.
+	require.Len(t, set.inputs, 1)
+
+	// Assert the wallet input is added using the set's deadline.
+	inp := set.inputs[0]
+	require.True(t, inp.params.DeadlineHeight.IsSome())
+	require.Equal(t, deadline, inp.params.DeadlineHeight.UnsafeFromSome())
+}
+
 // TestNeedWalletInput checks that NeedWalletInput correctly determines if a
 // wallet input is needed.
 func TestNeedWalletInput(t *testing.T) {