diff --git a/pkg/mocks/mockOpenOltClient.go b/pkg/mocks/mockOpenOltClient.go
index a59e3d0..5284bbd 100644
--- a/pkg/mocks/mockOpenOltClient.go
+++ b/pkg/mocks/mockOpenOltClient.go
@@ -31,7 +31,8 @@
 
 // MockOpenoltClient mock struct for OpenoltClient.
 type MockOpenoltClient struct {
-	counter int
+	counter     int
+	IsRestarted bool
 }
 
 // DisableOlt mocks the DisableOlt function of Openoltclient.
@@ -138,6 +139,7 @@
 func (ooc *MockOpenoltClient) Reboot(ctx context.Context, in *openolt.Empty, opts ...grpc.CallOption) (*openolt.Empty, error) {
 	if ooc.counter == 0 {
 		ooc.counter++
+		ooc.IsRestarted = true
 		return &openolt.Empty{}, nil
 	}
 	return nil, errors.New("reboot failed")