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/lnutils"
"github.com/lightningnetwork/lnd/lnwire"
"github.com/lightningnetwork/lnd/msgmux"
)
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
// 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
return fn.MapOptionZ(cfgMapper, func(mapper MsgMapper[Event]) bool {
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.
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
// this message.

View File

@ -13,6 +13,7 @@ import (
"github.com/lightningnetwork/lnd/chainntnfs"
"github.com/lightningnetwork/lnd/fn/v2"
"github.com/lightningnetwork/lnd/lnwire"
"github.com/lightningnetwork/lnd/msgmux"
"github.com/stretchr/testify/mock"
"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
// terminate the state machine.
wireError := &lnwire.Error{}
initMsg := &lnwire.Init{}
wireError := msgmux.PeerMsg{
Message: &lnwire.Error{},
}
initMsg := msgmux.PeerMsg{
Message: &lnwire.Init{},
}
dummyMapper.On("MapMsg", wireError).Return(
fn.Some(dummyEvents(&goToFin{})),
)
@ -456,7 +461,7 @@ func TestStateMachineMsgMapper(t *testing.T) {
// First, we'll verify that the CanHandle method works as expected.
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.
// We should transition to the final state.