diff --git a/extended/code_update/code_update_utils.robot b/extended/code_update/code_update_utils.robot
index b5bfe95..a762c03 100644
--- a/extended/code_update/code_update_utils.robot
+++ b/extended/code_update/code_update_utils.robot
@@ -11,6 +11,7 @@
 ${BMC_UPD_METHOD}    /org/openbmc/control/flash/bmc/action/update
 ${BMC_PREP_METHOD}   /org/openbmc/control/flash/bmc/action/PrepareForUpdate
 ${BMC_UPD_ATTR}      /org/openbmc/control/flash/bmc
+${HOST_SETTING}      /org/openbmc/settings/host0
 
 *** Keywords ***
 
@@ -97,3 +98,12 @@
     Open Connection And Log In
 
     Start Command   /sbin/reboot
+
+Set Policy Setting
+    [Documentation]   Set the given test policy
+    [arguments]   ${policy}
+
+    ${valueDict} =     create dictionary  data=${policy}
+    Write Attribute    ${HOST_SETTING}    power_policy   data=${valueDict}
+    ${currentPolicy}=  Read Attribute     ${HOST_SETTING}   power_policy
+    Should Be Equal    ${currentPolicy}   ${policy}
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}
diff --git a/lib/utils.robot b/lib/utils.robot
index 4f0acca..229b858 100644
--- a/lib/utils.robot
+++ b/lib/utils.robot
@@ -306,5 +306,5 @@
     [Arguments]       ${expected}
 
     ${current}=  Get BMC State
-    Should Contain  ${current}   ${expected}
+    Should Contain  ${expected}   ${current}