itest: wait for additional sweep if necessary

This fixes an issue where the lnd sweeper sometimes doesn't sweep all
timeout HTLCs in a single sweep TX. If that happens, we just need to
wait for an additional one.

Author: guggero

itest/litd_custom_channels_test.go

@@ -3,6 +3,7 @@ package itest
 import (
 	"bytes"
 	"context"
+	"encoding/hex"
 	"fmt"
 	"math"
 	"math/big"
@@ -12,6 +13,7 @@ import (
 	"github.com/btcsuite/btcd/btcec/v2/schnorr"
 	"github.com/btcsuite/btcd/btcutil"
 	"github.com/btcsuite/btcd/chaincfg/chainhash"
+	"github.com/btcsuite/btcd/wire"
 	"github.com/lightninglabs/taproot-assets/asset"
 	"github.com/lightninglabs/taproot-assets/itest"
 	"github.com/lightninglabs/taproot-assets/proof"
@@ -3261,8 +3263,8 @@ func runCustomChannelsHtlcForceClose(ctx context.Context, t *harnessTest,
 	// At this point, both sides should have 4 (or +4 with MPP) HTLCs
 	// active.
 	numHtlcs := 4
+	numAdditionalShards := assetInvoiceAmt / assetsPerMPPShard
 	if mpp {
-		numAdditionalShards := assetInvoiceAmt / assetsPerMPPShard
 		numHtlcs += numAdditionalShards * 2
 	}
 	t.Logf("Asserting both Alice and Bob have %d HTLCs...", numHtlcs)
@@ -3521,6 +3523,10 @@ func runCustomChannelsHtlcForceClose(ctx context.Context, t *harnessTest,
 
 	// We'll wait for both Alice and Bob to present their respective sweeps
 	// to the sweeper.
+	numTimeoutHTLCs := 1
+	if mpp {
+		numTimeoutHTLCs += numAdditionalShards
+	}
 	assertSweepExists(
 		t.t, alice,
 		walletrpc.WitnessType_TAPROOT_HTLC_LOCAL_OFFERED_TIMEOUT,
@@ -3542,6 +3548,58 @@ func runCustomChannelsHtlcForceClose(ctx context.Context, t *harnessTest,
 	// Finally, we'll mine a single block to confirm them.
 	mineBlocks(t, net, 1, 2)
 
+	// Make sure Bob swept all his HTLCs.
+	bobSweeps, err := bob.WalletKitClient.ListSweeps(
+		ctx, &walletrpc.ListSweepsRequest{
+			Verbose: true,
+		},
+	)
+	require.NoError(t.t, err)
+
+	var bobSweepTx *wire.MsgTx
+	for _, sweep := range bobSweeps.GetTransactionDetails().Transactions {
+		for _, tx := range timeoutSweeps {
+			if sweep.TxHash == tx.String() {
+				txBytes, err := hex.DecodeString(sweep.RawTxHex)
+				require.NoError(t.t, err)
+
+				bobSweepTx = &wire.MsgTx{}
+				err = bobSweepTx.Deserialize(
+					bytes.NewReader(txBytes),
+				)
+				require.NoError(t.t, err)
+			}
+		}
+	}
+	require.NotNil(t.t, bobSweepTx, "Bob's sweep transaction not found")
+
+	// There's always an extra input that pays for the fees. So we can only
+	// count the remainder as HTLC inputs.
+	numSweptHTLCs := len(bobSweepTx.TxIn) - 1
+
+	// If we didn't yet sweep all HTLCs, then we need to wait for another
+	// sweep.
+	if numSweptHTLCs < numTimeoutHTLCs {
+		assertSweepExists(
+			t.t, bob,
+			// nolint: lll
+			walletrpc.WitnessType_TAPROOT_HTLC_OFFERED_REMOTE_TIMEOUT,
+		)
+
+		t.Logf("Confirming additional HTLC timeout sweep txns")
+
+		additionalTimeoutSweeps, err := waitForNTxsInMempool(
+			net.Miner.Client, 1, shortTimeout,
+		)
+		require.NoError(t.t, err)
+
+		t.Logf("Asserting balance on additional timeout sweeps: %v",
+			additionalTimeoutSweeps)
+
+		// Finally, we'll mine a single block to confirm them.
+		mineBlocks(t, net, 1, 1)
+	}
+
 	// At this point, Bob's balance should be incremented by an additional
 	// HTLC value.
 	bobExpectedBalance += uint64(assetInvoiceAmt - 1)