diff --git a/adaptercore/olt_state_transitions_test.go b/adaptercore/olt_state_transitions_test.go
index 31b7d56..6c256a1 100644
--- a/adaptercore/olt_state_transitions_test.go
+++ b/adaptercore/olt_state_transitions_test.go
@@ -17,9 +17,11 @@
 package adaptercore
 
 import (
+	"context"
 	"errors"
 	"reflect"
 	"testing"
+	"time"
 )
 
 /**
@@ -43,9 +45,9 @@
 	transition := Transition{
 		previousState: []DeviceState{deviceStateConnected},
 		currentState:  deviceStateConnected,
-		after: []TransitionHandler{func() error {
+		after: []TransitionHandler{func(ctx context.Context) error {
 			return nil
-		}, func() error {
+		}, func(ctx context.Context) error {
 			return errors.New("transition error")
 		}},
 	}
@@ -61,9 +63,9 @@
 	transition := Transition{
 		previousState: []DeviceState{deviceStateConnected},
 		currentState:  deviceStateConnected,
-		before: []TransitionHandler{func() error {
+		before: []TransitionHandler{func(ctx context.Context) error {
 			return nil
-		}, func() error {
+		}, func(ctx context.Context) error {
 			return errors.New("transition error")
 		}},
 	}
@@ -121,7 +123,9 @@
 				transitions:        tt.fields.transitions,
 				currentDeviceState: tt.fields.currentDeviceState,
 			}
-			tMap.Handle(tt.args.trigger)
+			ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+			defer cancel()
+			tMap.Handle(ctx, tt.args.trigger)
 		})
 	}
 }