lnwallet: add tests for the new rebroadcaster logic

We needed to copy some mocks from elsewhere in the codebase, as
otherwise we'd run into an import cycle.

lnwallet/rebroadcaster_test.go

@@ -0,0 +1,197 @@
+package lnwallet
+
+import (
+	"sync/atomic"
+	"testing"
+	"time"
+
+	"github.com/btcsuite/btcd/chaincfg/chainhash"
+	"github.com/btcsuite/btcd/wire"
+
+	"github.com/lightningnetwork/lnd/chainntnfs"
+	"github.com/lightningnetwork/lnd/lnutils"
+	"github.com/stretchr/testify/require"
+)
+
+type mockRebroadcaster struct {
+	started atomic.Bool
+
+	rebroadcastAttempt chan struct{}
+
+	falseStart bool
+
+	confSignal chan struct{}
+
+	startSignal chan struct{}
+}
+
+func newMockRebroadcaster(falseStart bool) *mockRebroadcaster {
+	return &mockRebroadcaster{
+		rebroadcastAttempt: make(chan struct{}, 1),
+		falseStart:         falseStart,
+		confSignal:         make(chan struct{}, 1),
+		startSignal:        make(chan struct{}),
+	}
+}
+
+func (m *mockRebroadcaster) Start() error {
+	if !m.falseStart {
+		defer m.started.Store(true)
+		defer close(m.startSignal)
+	}
+
+	return nil
+}
+
+func (m *mockRebroadcaster) Started() bool {
+	return m.started.Load()
+}
+
+func (m *mockRebroadcaster) Stop() {
+}
+
+// Broadcast enqueues a transaction to be rebroadcast until it's been
+// confirmed.
+func (m *mockRebroadcaster) Broadcast(tx *wire.MsgTx) error {
+	m.rebroadcastAttempt <- struct{}{}
+	return nil
+}
+
+func (m *mockRebroadcaster) MarkAsConfirmed(txid chainhash.Hash) {
+	m.confSignal <- struct{}{}
+}
+
+func assertBroadcasterBypass(t *testing.T, wallet *LightningWallet,
+	rebroadcaster *mockRebroadcaster,
+	walletController *mockWalletController) {
+
+	testTx := wire.NewMsgTx(2)
+	timeout := time.Second * 1
+
+	require.NoError(t, wallet.PublishTransaction(testTx, ""))
+
+	// The tx should go to the backend.
+	_, err := lnutils.RecvOrTimeout(
+		walletController.PublishedTransactions, timeout,
+	)
+	require.NoError(t, err)
+
+	// It shouldn't go to the rebroadcaster.
+	select {
+	case <-time.After(timeout):
+	case <-rebroadcaster.rebroadcastAttempt:
+		t.Fatal("tx sent to rebroadcaster")
+	}
+}
+
+func assertBroadcasterSend(t *testing.T, wallet *LightningWallet,
+	rebroadcaster *mockRebroadcaster,
+	walletController *mockWalletController) {
+
+	testTx := wire.NewMsgTx(2)
+	testTx.AddTxOut(&wire.TxOut{})
+
+	timeout := time.Second * 1
+
+	require.NoError(t, wallet.PublishTransaction(testTx, ""))
+
+	// The tx should go to the backend.
+	_, err := lnutils.RecvOrTimeout(
+		walletController.PublishedTransactions, timeout,
+	)
+	require.NoError(t, err)
+
+	// It should also go to the rebroadcaster.
+	select {
+	case <-time.After(timeout):
+		t.Fatal("tx not sent to rebroadcaster")
+	case <-rebroadcaster.rebroadcastAttempt:
+	}
+}
+
+// TestWalletRebroadcaster tests that the wallet properly manages the existence
+// or lack of existence of the rebroadcaster, and also properly marks the
+// transaction as confirmed.
+func TestWalletRebroadcaster(t *testing.T) {
+	t.Parallel()
+
+	rebroadcaster := newMockRebroadcaster(false)
+	walletController := &mockWalletController{
+		PublishedTransactions: make(chan *wire.MsgTx, 1),
+	}
+	chainIO := &mockChainIO{}
+	notifier := &mockChainNotifier{
+		SpendChan: make(chan *chainntnfs.SpendDetail, 1),
+		EpochChan: make(chan *chainntnfs.BlockEpoch, 1),
+		ConfChan:  make(chan *chainntnfs.TxConfirmation, 1),
+	}
+	cfg := &Config{
+		Rebroadcaster:    rebroadcaster,
+		WalletController: walletController,
+		Notifier:         notifier,
+		ChainIO:          chainIO,
+	}
+
+	t.Run("rebroadcast bypass", func(t *testing.T) {
+		// We'll make a copy of the config, but without the
+		// broadcaster.
+		testCfg := *cfg
+		testCfg.Rebroadcaster = nil
+
+		wallet, err := NewLightningWallet(testCfg)
+		require.NoError(t, err)
+		require.NoError(t, wallet.Startup())
+
+		// If we try to broadcast, it should go straight to the wallet
+		// backend and skip the broadcaster.
+		assertBroadcasterBypass(
+			t, wallet, rebroadcaster, walletController,
+		)
+
+		wallet.Shutdown()
+
+		// If we make a new wallet, that has the broadcaster, but
+		// hasn't started yet, we should see the same behavior.
+		testCfg.Rebroadcaster = newMockRebroadcaster(true)
+
+		wallet, err = NewLightningWallet(testCfg)
+		require.NoError(t, err)
+		require.NoError(t, wallet.Startup())
+
+		assertBroadcasterBypass(
+			t, wallet, rebroadcaster, walletController,
+		)
+
+		wallet.Shutdown()
+	})
+
+	t.Run("rebroadcast normal", func(t *testing.T) {
+		wallet, err := NewLightningWallet(*cfg)
+		require.NoError(t, err)
+		require.NoError(t, wallet.Startup())
+
+		defer wallet.Shutdown()
+
+		// Wait for the broadcaster to start.
+		_, err = lnutils.RecvOrTimeout(
+			rebroadcaster.startSignal, time.Second,
+		)
+		require.NoError(t, err)
+
+		// We'll now broadcast a new test transaction, asserting that
+		// it goes to both the backend and the rebroadcaster.
+		assertBroadcasterSend(
+			t, wallet, rebroadcaster, walletController,
+		)
+
+		// We'll now mark the transaction as confirmed, and assert that
+		// the rebroadcaster was notified.
+		notifier.ConfChan <- &chainntnfs.TxConfirmation{}
+
+		_, err = lnutils.RecvOrTimeout(
+			rebroadcaster.confSignal, time.Second,
+		)
+		require.NoError(t, err)
+
+	})
+}