diff --git a/libraries/voltctl.robot b/libraries/voltctl.robot
index f45518c..57ec5c8 100755
--- a/libraries/voltctl.robot
+++ b/libraries/voltctl.robot
@@ -702,12 +702,12 @@
     List Should Not Contain Value    ${ids}    ${id}
 
 Reboot ONU
-    [Arguments]    ${onu_id}
+    [Arguments]    ${onu_id}    ${validate_device}=True
     [Documentation]   Using voltctl command reboot ONU and verify that ONU comes up to running state
     ${rc}    ${devices}=    Run and Return Rc and Output    ${VOLTCTL_CONFIG}; voltctl device reboot ${onu_id}
     Should Be Equal As Integers    ${rc}    0
-    Run Keyword And Continue On Failure    Wait Until Keyword Succeeds    60s   1s    Validate ONU Device By Device Id
-    ...    ENABLED    DISCOVERED    REACHABLE    rebooting   ${onu_id}
+    Run Keyword If    ${validate_device}    Run Keyword And Continue On Failure    Wait Until Keyword Succeeds
+    ...    60s   1s    Validate ONU Device By Device Id    ENABLED    DISCOVERED    REACHABLE    rebooting   ${onu_id}
 
 Assert ONUs in Voltha
     [Arguments]    ${count}
diff --git a/tests/openonu-go-adapter/Voltha_ONUStateTests.robot b/tests/openonu-go-adapter/Voltha_ONUStateTests.robot
index a493133..999e34e 100755
--- a/tests/openonu-go-adapter/Voltha_ONUStateTests.robot
+++ b/tests/openonu-go-adapter/Voltha_ONUStateTests.robot
@@ -400,10 +400,11 @@
     FOR    ${I}    IN RANGE    0    ${num_onus}
         ${src}=    Set Variable    ${hosts.src[${I}]}
         ${onu_device_id}=    Get Device ID From SN    ${src['onu']}
-        Reboot ONU    ${onu_device_id}
+        Reboot ONU    ${onu_device_id}   False
     END
-    Run Keyword If    ${has_dataplane}    Sleep    30s
-    Do Disable Enable Onu Test    checkstatebeforedisable=False
+    Run Keyword Unless    ${has_dataplane}    Do Current State Test All Onus    rebooting
+    Run Keyword Unless    ${has_dataplane}    Do Disable Enable Onu Test    checkstatebeforedisable=False
+    Run Keyword If    ${has_dataplane}    Do Current State Test All Onus    omci-flows-pushed
     Do Onu Port Check
 
 Do Disable Onu Device
@@ -496,15 +497,16 @@
     [Arguments]    ${state}
     # create state lists with corresponding return values
     #                             ADMIN-STATE OPER-STATUS   CONNECT-STATUS ONU-STATE (number/name)
-    ${state1}    Create List      ENABLED     ACTIVATING    REACHABLE      1    activating-onu
-    ${state2}    Create List      ENABLED     ACTIVATING    REACHABLE      2    starting-openomci
-    ${state3}    Create List      ENABLED     ACTIVATING    REACHABLE      3    discovery-mibsync-complete
-    ${state4}    Create List      ENABLED     ACTIVE        REACHABLE      4    initial-mib-downloaded
-    ${state5}    Create List      ENABLED     ACTIVE        REACHABLE      5    tech-profile-config-download-success
-    ${state6}    Create List      ENABLED     ACTIVE        REACHABLE      6    omci-flows-pushed
-    ${state7}    Create List      DISABLED    UNKNOWN       REACHABLE      7    omci-admin-lock
-    ${state8}    Create List      ENABLED     ACTIVE        REACHABLE      8    onu-reenabled
-    ${state9}    Create List      ENABLED     DISCOVERED    UNREACHABLE    9    stopping-openomci
+    ${state1}     Create List      ENABLED     ACTIVATING    REACHABLE       1    activating-onu
+    ${state2}     Create List      ENABLED     ACTIVATING    REACHABLE       2    starting-openomci
+    ${state3}     Create List      ENABLED     ACTIVATING    REACHABLE       3    discovery-mibsync-complete
+    ${state4}     Create List      ENABLED     ACTIVE        REACHABLE       4    initial-mib-downloaded
+    ${state5}     Create List      ENABLED     ACTIVE        REACHABLE       5    tech-profile-config-download-success
+    ${state6}     Create List      ENABLED     ACTIVE        REACHABLE       6    omci-flows-pushed
+    ${state7}     Create List      DISABLED    UNKNOWN       REACHABLE       7    omci-admin-lock
+    ${state8}     Create List      ENABLED     ACTIVE        REACHABLE       8    onu-reenabled
+    ${state9}     Create List      ENABLED     DISCOVERED    UNREACHABLE     9    stopping-openomci
+    ${state10}    Create List      ENABLED     DISCOVERED    REACHABLE      10    rebooting
     ${admin_state}    ${oper_status}    ${connect_status}    ${onu_state_nb}    ${onu_state}=    Set Variable If
     ...    '${state}'=='1' or '${state}'=='activating-onu'                          ${state1}
     ...    '${state}'=='2' or '${state}'=='starting-openomci'                       ${state2}
@@ -515,4 +517,5 @@
     ...    '${state}'=='7' or '${state}'=='omci-admin-lock'                         ${state7}
     ...    '${state}'=='8' or '${state}'=='onu-reenabled'                           ${state8}
     ...    '${state}'=='9' or '${state}'=='stopping-openomci'                       ${state9}
+    ...    '${state}'=='10' or '${state}'=='rebooting'                              ${state10}
     [Return]    ${admin_state}    ${oper_status}    ${connect_status}    ${onu_state_nb}    ${onu_state}