diff --git a/extended/code_update/update_bmc.robot b/extended/code_update/update_bmc.robot
index 461bddb..b4e44e3 100644
--- a/extended/code_update/update_bmc.robot
+++ b/extended/code_update/update_bmc.robot
@@ -8,6 +8,9 @@
 ...               Update work flow sequence:
 ...                 - User input BMC File existence check
 ...                 - Ping Test and REST authentication
+...                 - Set Host Power host setting Policy to RESTORE_LAST_STATE
+...                   On reboot this policy would ensure the BMC comes
+...                   online and stays at HOST_POWERED_OFF state.
 ...                 - Issue poweroff
 ...                 - Prepare for Update
 ...                 - Wait for BMC to come online clean
@@ -20,6 +23,7 @@
 ...                 - Version check post update
 
 Resource          code_update_utils.robot
+Resource          ../../lib/boot/boot_resource_master.robot
 
 *** Variables ***
 
@@ -39,14 +43,16 @@
     ...     Pass Execution   Same Driver version installed
 
     Initiate Power Off
+    Set Policy Setting   RESTORE_LAST_STATE
     Prepare For Update
 
     # Wait time is increased temporary to 10 mins due
     # to openbmc/openbmc#673
     Check If BMC is Up    10 min   10 sec
 
+    @{states}=   Create List   BMC_READY   HOST_POWERED_OFF
     Wait Until Keyword Succeeds
-    ...    10 min   10 sec   Verify BMC State   BMC_READY
+    ...    10 min   10 sec   Verify BMC State   ${states}
 
     Preserve BMC Network Setting
     SCP Tar Image File to BMC   ${FILE_PATH}