protofsm: update StateMachine to meet new msgmux.Endpoint interface

This commit is contained in:
Olaoluwa Osuntokun 2025-02-27 13:56:38 -08:00
parent bf3007a2ce
commit ed8a672bd3
2 changed files with 11 additions and 5 deletions

View File

@ -14,6 +14,7 @@ import (
"github.com/lightningnetwork/lnd/fn/v2" "github.com/lightningnetwork/lnd/fn/v2"
"github.com/lightningnetwork/lnd/lnutils" "github.com/lightningnetwork/lnd/lnutils"
"github.com/lightningnetwork/lnd/lnwire" "github.com/lightningnetwork/lnd/lnwire"
"github.com/lightningnetwork/lnd/msgmux"
) )
const ( const (
@ -250,7 +251,7 @@ func (s *StateMachine[Event, Env]) SendEvent(ctx context.Context, event Event) {
// CanHandle returns true if the target message can be routed to the state // CanHandle returns true if the target message can be routed to the state
// machine. // machine.
func (s *StateMachine[Event, Env]) CanHandle(msg lnwire.Message) bool { func (s *StateMachine[Event, Env]) CanHandle(msg msgmux.PeerMsg) bool {
cfgMapper := s.cfg.MsgMapper cfgMapper := s.cfg.MsgMapper
return fn.MapOptionZ(cfgMapper, func(mapper MsgMapper[Event]) bool { return fn.MapOptionZ(cfgMapper, func(mapper MsgMapper[Event]) bool {
return mapper.MapMsg(msg).IsSome() return mapper.MapMsg(msg).IsSome()
@ -267,7 +268,7 @@ func (s *StateMachine[Event, Env]) Name() string {
// returned indicating that the message was processed. Otherwise, false is // returned indicating that the message was processed. Otherwise, false is
// returned. // returned.
func (s *StateMachine[Event, Env]) SendMessage(ctx context.Context, func (s *StateMachine[Event, Env]) SendMessage(ctx context.Context,
msg lnwire.Message) bool { msg msgmux.PeerMsg) bool {
// If we have no message mapper, then return false as we can't process // If we have no message mapper, then return false as we can't process
// this message. // this message.

View File

@ -13,6 +13,7 @@ import (
"github.com/lightningnetwork/lnd/chainntnfs" "github.com/lightningnetwork/lnd/chainntnfs"
"github.com/lightningnetwork/lnd/fn/v2" "github.com/lightningnetwork/lnd/fn/v2"
"github.com/lightningnetwork/lnd/lnwire" "github.com/lightningnetwork/lnd/lnwire"
"github.com/lightningnetwork/lnd/msgmux"
"github.com/stretchr/testify/mock" "github.com/stretchr/testify/mock"
"github.com/stretchr/testify/require" "github.com/stretchr/testify/require"
) )
@ -429,8 +430,12 @@ func TestStateMachineMsgMapper(t *testing.T) {
// The only thing we know how to map is the error message, which'll // The only thing we know how to map is the error message, which'll
// terminate the state machine. // terminate the state machine.
wireError := &lnwire.Error{} wireError := msgmux.PeerMsg{
initMsg := &lnwire.Init{} Message: &lnwire.Error{},
}
initMsg := msgmux.PeerMsg{
Message: &lnwire.Init{},
}
dummyMapper.On("MapMsg", wireError).Return( dummyMapper.On("MapMsg", wireError).Return(
fn.Some(dummyEvents(&goToFin{})), fn.Some(dummyEvents(&goToFin{})),
) )
@ -456,7 +461,7 @@ func TestStateMachineMsgMapper(t *testing.T) {
// First, we'll verify that the CanHandle method works as expected. // First, we'll verify that the CanHandle method works as expected.
require.True(t, stateMachine.CanHandle(wireError)) require.True(t, stateMachine.CanHandle(wireError))
require.False(t, stateMachine.CanHandle(&lnwire.Init{})) require.False(t, stateMachine.CanHandle(initMsg))
// Next, we'll attempt to send the wire message into the state machine. // Next, we'll attempt to send the wire message into the state machine.
// We should transition to the final state. // We should transition to the final state.