COCOS-108 - Enhanced state machine Synchronization (#112)

* Enhanced state machine synchronization

Introduced a wait group to the state machine to manage go-routine synchronization effectively. The Start method now ensures that the state machine's execution is tracked, preventing potential early termination of go-routines during state transitions. The test suite is updated to wait for the go-routine to complete before proceeding, ensuring more reliable test execution.

Signed-off-by: SammyOina <sammyoina@gmail.com>

* Remove unnecessary code and fix test case in agent package

Signed-off-by: SammyOina <sammyoina@gmail.com>

---------

Signed-off-by: SammyOina <sammyoina@gmail.com>

agent/state.go

@@ -41,6 +41,7 @@ type StateMachine struct {
 	Transitions    map[state]map[event]state
 	StateFunctions map[state]func()
 	logger         *slog.Logger
+	wg             *sync.WaitGroup
 }
 
 // NewStateMachine creates a new StateMachine.
@@ -51,6 +52,7 @@ func NewStateMachine(logger *slog.Logger) *StateMachine {
 		Transitions:    make(map[state]map[event]state),
 		StateFunctions: make(map[state]func()),
 		logger:         logger,
+		wg:             &sync.WaitGroup{},
 	}
 
 	sm.Transitions[idle] = make(map[event]state)
@@ -76,6 +78,8 @@ func NewStateMachine(logger *slog.Logger) *StateMachine {
 
 // Start the state machine.
 func (sm *StateMachine) Start(ctx context.Context) {
+	sm.wg.Add(1)
+	defer sm.wg.Done()
 	for {
 		select {
 		case event := <-sm.EventChan:

agent/state_test.go

@@ -27,12 +27,11 @@ func TestStateMachineTransitions(t *testing.T) {
 	for _, testCase := range testCases {
 		t.Run(fmt.Sprintf("Transition from %v to %v", testCase.fromState, testCase.expected), func(t *testing.T) {
 			sm := NewStateMachine(mglog.NewMock())
-			done := make(chan struct{})
 			ctx, cancel := context.WithCancel(context.Background())
 			go func() {
 				sm.Start(ctx)
-				close(done)
 			}()
+			sm.wg.Wait()
 			sm.SetState(testCase.fromState)
 
 			sm.SendEvent(testCase.event)
@@ -42,7 +41,6 @@ func TestStateMachineTransitions(t *testing.T) {
 			}
 			close(sm.EventChan)
 			cancel()
-			<-done
 		})
 	}
 }