diff --git a/protofsm/state_machine_test.go b/protofsm/state_machine_test.go index 0432f386b7..fc30fcefc3 100644 --- a/protofsm/state_machine_test.go +++ b/protofsm/state_machine_test.go @@ -251,14 +251,14 @@ func TestStateMachineOnInitDaemonEvent(t *testing.T) { // message adapter is called on start up. adapters.On("SendMessages", *pub1, mock.Anything).Return(nil) - stateMachine.Start() - defer stateMachine.Stop() - // As we're triggering internal events, we'll also subscribe to the set // of new states so we can assert as we go. stateSub := stateMachine.RegisterStateEvents() defer stateMachine.RemoveStateSub(stateSub) + stateMachine.Start() + defer stateMachine.Stop() + // Assert that we go from the starting state to the final state. The // state machine should now also be on the final terminal state. expectedStates := []State[dummyEvents, *dummyEnv]{ @@ -292,14 +292,15 @@ func TestStateMachineInternalEvents(t *testing.T) { InitEvent: fn.None[DaemonEvent](), } stateMachine := NewStateMachine(cfg) - stateMachine.Start() - defer stateMachine.Stop() // As we're triggering internal events, we'll also subscribe to the set // of new states so we can assert as we go. stateSub := stateMachine.RegisterStateEvents() defer stateMachine.RemoveStateSub(stateSub) + stateMachine.Start() + defer stateMachine.Stop() + // For this transition, we'll send in the emitInternal event, which'll // send us back to the starting event, but emit an internal event. stateMachine.SendEvent(&emitInternal{}) @@ -343,14 +344,15 @@ func TestStateMachineDaemonEvents(t *testing.T) { InitEvent: fn.None[DaemonEvent](), } stateMachine := NewStateMachine(cfg) - stateMachine.Start() - defer stateMachine.Stop() // As we're triggering internal events, we'll also subscribe to the set // of new states so we can assert as we go. stateSub := stateMachine.RegisterStateEvents() defer stateMachine.RemoveStateSub(stateSub) + stateMachine.Start() + defer stateMachine.Stop() + // As soon as we send in the daemon event, we expect the // disable+broadcast events to be processed, as they are unconditional. adapters.On( @@ -428,14 +430,17 @@ func TestStateMachineMsgMapper(t *testing.T) { MsgMapper: fn.Some[MsgMapper[dummyEvents]](dummyMapper), } stateMachine := NewStateMachine(cfg) - stateMachine.Start() - defer stateMachine.Stop() // As we're triggering internal events, we'll also subscribe to the set // of new states so we can assert as we go. + // + // We register before calling Start to ensure we don't miss any events. stateSub := stateMachine.RegisterStateEvents() defer stateMachine.RemoveStateSub(stateSub) + stateMachine.Start() + defer stateMachine.Stop() + // First, we'll verify that the CanHandle method works as expected. require.True(t, stateMachine.CanHandle(wireError)) require.False(t, stateMachine.CanHandle(&lnwire.Init{}))