Merge pull request #9726 from Roasbeef/protofsm-full-block-conf

protofsm: add option to allow conf resp to return full block

Author: Roasbeef

protofsm/daemon_events.go

@@ -116,6 +116,11 @@ type RegisterConf[Event any] struct {
 	// transaction needs to dispatch an event.
 	NumConfs fn.Option[uint32]
 
+	// FullBlock is a boolean that indicates whether we want the full block
+	// in the returned response. This is useful if callers want to create an
+	// SPV proof for the transaction post conf.
+	FullBlock bool
+
 	// PostConfMapper is a special conf mapper, that if present, will be
 	// used to map the protofsm confirmation event to a custom event.
 	PostConfMapper fn.Option[ConfMapper[Event]]

protofsm/state_machine.go

@@ -510,10 +510,15 @@ func (s *StateMachine[Event, Env]) executeDaemonEvent(ctx context.Context,
 		s.log.DebugS(ctx, "Registering conf",
 			"txid", daemonEvent.Txid)
 
+		var opts []chainntnfs.NotifierOption
+		if daemonEvent.FullBlock {
+			opts = append(opts, chainntnfs.WithIncludeBlock())
+		}
+
 		numConfs := daemonEvent.NumConfs.UnwrapOr(1)
 		confEvent, err := s.cfg.Daemon.RegisterConfirmationsNtfn(
 			&daemonEvent.Txid, daemonEvent.PkScript,
-			numConfs, daemonEvent.HeightHint,
+			numConfs, daemonEvent.HeightHint, opts...,
 		)
 		if err != nil {
 			return fmt.Errorf("unable to register conf: %w", err)

protofsm/state_machine_test.go

@@ -49,6 +49,7 @@ func (c *confDetailsEvent) dummy() {
 }
 
 type registerConf struct {
+	fullBlock bool
 }
 
 func (r *registerConf) dummy() {
@@ -173,6 +174,7 @@ func (d *dummyStateStart) ProcessEvent(event dummyEvents, env *dummyEnv,
 			Txid:       chainhash.Hash{1},
 			PkScript:   []byte{0x01},
 			HeightHint: 100,
+			FullBlock:  newEvent.fullBlock,
 			PostConfMapper: fn.Some[ConfMapper[dummyEvents]](
 				confMapper,
 			),
@@ -383,7 +385,8 @@ func (d *dummyAdapters) RegisterConfirmationsNtfn(txid *chainhash.Hash,
 	opts ...chainntnfs.NotifierOption,
 ) (*chainntnfs.ConfirmationEvent, error) {
 
-	args := d.Called(txid, pkScript, numConfs)
+	// Pass opts as the last argument to the mock call checker.
+	args := d.Called(txid, pkScript, numConfs, heightHint, opts)
 
 	err := args.Error(0)
 
@@ -589,11 +592,12 @@ func TestStateMachineDaemonEvents(t *testing.T) {
 	env.AssertExpectations(t)
 }
 
-// TestStateMachineConfMapper tests that the state machine is able to properly
-// map the confirmation event into a custom event that can be used to trigger a
-// state transition.
-func TestStateMachineConfMapper(t *testing.T) {
-	t.Parallel()
+// testStateMachineConfMapperImpl is a helper function that encapsulates the
+// core logic for testing the confirmation mapping functionality of the state
+// machine. It takes a boolean flag `fullBlock` to determine whether to test the
+// scenario where full block details are requested in the confirmation
+// notification.
+func testStateMachineConfMapperImpl(t *testing.T, fullBlock bool) {
 	ctx := context.Background()
 
 	// Create the state machine.
@@ -614,35 +618,69 @@ func TestStateMachineConfMapper(t *testing.T) {
 	stateMachine.Start(ctx)
 	defer stateMachine.Stop()
 
-	// Expect the RegisterConfirmationsNtfn call when we send the event.
-	// We use NumConfs=1 as the default.
-	adapters.On(
-		"RegisterConfirmationsNtfn", &chainhash.Hash{1}, []byte{0x01},
-		uint32(1),
-	).Return(nil)
+	// Define the expected arguments for the mock call.
+	expectedTxid := &chainhash.Hash{1}
+	expectedPkScript := []byte{0x01}
+	expectedNumConfs := uint32(1)
+	expectedHeightHint := uint32(100)
+
+	// Set up the mock expectation based on the FullBlock flag. We use
+	// mock.MatchedBy to assert the options passed.
+	if fullBlock {
+		// Expect WithIncludeBlock() option when FullBlock is true.
+		adapters.On(
+			"RegisterConfirmationsNtfn",
+			expectedTxid, expectedPkScript,
+			expectedNumConfs, expectedHeightHint,
+			mock.MatchedBy(
+				func(opts []chainntnfs.NotifierOption) bool {
+					// Check if exactly one option is passed
+					// and it's the correct type. Unless we
+					// use reflect, we can introspect into
+					// the private fields.
+					return len(opts) == 1
+				},
+			),
+		).Return(nil)
+	} else {
+		// Expect no options when FullBlock is false.
+		adapters.On(
+			"RegisterConfirmationsNtfn",
+			expectedTxid, expectedPkScript,
+			expectedNumConfs, expectedHeightHint,
+			mock.MatchedBy(func(opts []chainntnfs.NotifierOption) bool { //nolint:ll
+				return len(opts) == 0
+			}),
+		).Return(nil)
+	}
+
+	// Create the registerConf event with the specified FullBlock value.
+	regConfEvent := &registerConf{
+		fullBlock: fullBlock,
+	}
 
 	// Send the event that triggers RegisterConf emission.
-	stateMachine.SendEvent(ctx, &registerConf{})
+	stateMachine.SendEvent(ctx, regConfEvent)
 
 	// We should transition back to the starting state initially.
 	expectedStates := []State[dummyEvents, *dummyEnv]{
 		&dummyStateStart{}, &dummyStateStart{},
 	}
 	assertStateTransitions(t, stateSub, expectedStates)
 
-	// Assert the registration call was made.
+	// Assert the registration call was made with the correct arguments
+	// (including options).
 	adapters.AssertExpectations(t)
 
 	// Now, simulate the confirmation event coming back from the notifier.
-	// Populate it with some data to be mapped.
 	simulatedConf := &chainntnfs.TxConfirmation{
 		BlockHash:   &chainhash.Hash{2},
 		BlockHeight: 123,
 	}
 	adapters.confChan <- simulatedConf
 
 	// This should trigger the mapper and send the confDetailsEvent,
-	// transitioning us to the final state.
+	// transitioning us to the confirmed state.
 	expectedStates = []State[dummyEvents, *dummyEnv]{&dummyStateConfirmed{}}
 	assertStateTransitions(t, stateSub, expectedStates)
 
@@ -662,6 +700,23 @@ func TestStateMachineConfMapper(t *testing.T) {
 	env.AssertExpectations(t)
 }
 
+// TestStateMachineConfMapper tests the confirmation mapping functionality using
+// subtests driven by the testStateMachineConfMapperImpl helper function. It
+// covers scenarios both with and without requesting the full block details.
+func TestStateMachineConfMapper(t *testing.T) {
+	t.Parallel()
+
+	t.Run("full block false", func(t *testing.T) {
+		t.Parallel()
+		testStateMachineConfMapperImpl(t, false)
+	})
+
+	t.Run("full block true", func(t *testing.T) {
+		t.Parallel()
+		testStateMachineConfMapperImpl(t, true)
+	})
+}
+
 // TestStateMachineSpendMapper tests that the state machine is able to properly
 // map the spend event into a custom event that can be used to trigger a state
 // transition.