diff --git a/extended/code_update/code_update.robot b/extended/code_update/code_update.robot
index 1d3304f..0d57788 100644
--- a/extended/code_update/code_update.robot
+++ b/extended/code_update/code_update.robot
@@ -23,6 +23,8 @@
 Resource          ../lib/openbmc_ffdc.robot
 Resource          ../../lib/code_update_utils.robot
 Resource          ../../lib/boot_utils.robot
+Resource          ../../lib/utils.robot
+Resource          code_update_utils.robot
 
 Test Teardown     Code Update Teardown
 
@@ -70,7 +72,10 @@
     ${software_state}=  Read Properties  ${SOFTWARE_VERSION_URI}${version_id}
     Should Be Equal As Strings  &{software_state}[Activation]  ${ACTIVE}
 
-    OBMC Reboot (off)
+    # TODO: Switch OBMC Reboot (off) once it's fixed
+    Trigger Warm Reset Via Reboot
+    Check If BMC is Up
+    Wait For BMC Ready
 
 
 Post Update Boot To OS