diff --git a/extended/code_update/update_bmc.robot b/extended/code_update/update_bmc.robot
index b82ae13..a9a0968 100644
--- a/extended/code_update/update_bmc.robot
+++ b/extended/code_update/update_bmc.robot
@@ -144,7 +144,11 @@
     Check Boot Count And Time
     Prune Journal Log
     Power Off Request
-    Set Policy Setting   RESTORE_LAST_STATE
+    ${status}=  Run Keyword And Return Status  Set BMC Power Policy
+    Run Keyword If  '${status}' == '${False}'
+    ...  Set BMC Power Policy
+    ...  ${HOST_SETTING}  power_policy  RESTORE_LAST_STATE
+
     Prepare For Update
     Check If BMC is Up  20 min  10 sec
     Check Boot Count And Time
diff --git a/lib/obmc_boot_test.py b/lib/obmc_boot_test.py
index 07e9f77..76f804d 100755
--- a/lib/obmc_boot_test.py
+++ b/lib/obmc_boot_test.py
@@ -350,7 +350,9 @@
 
     grp.rqprint_pgm_header()
 
-    grk.run_key("Set BMC Power Policy  RESTORE_LAST_STATE")
+    # TODO: Need to fix to switch between old and new.
+    # Disabling for a brief moment.
+    #grk.run_key("Set BMC Power Policy  ${RESTORE_LAST_STATE}")
 
     initial_plug_in_setup()
 
diff --git a/lib/utils.robot b/lib/utils.robot
index d9fa770..bc10fd8 100755
--- a/lib/utils.robot
+++ b/lib/utils.robot
@@ -724,15 +724,21 @@
 
 Set BMC Power Policy
     [Documentation]   Set the given BMC power policy.
-    [arguments]   ${policy}
+    [arguments]  ${power_restore_uri}=${POWER_RESTORE_URI}
+    ...          ${policy_attribute}=PowerRestorePolicy
+    ...          ${policy}=${RESTORE_LAST_STATE}
 
     # Description of argument(s):
-    # policy    Power restore policy.
+    # power_restore_uri    Power restore policy url.
+    #                      By default "/xyz/openbmc_project/control/host0/power_restore_policy/".
+    # policy               Power restore policy.
+    #                      By default ""xyz.openbmc_project.Control.Power.RestorePolicy.Policy.Restore".
 
     ${valueDict}=  Create Dictionary  data=${policy}
     Write Attribute
-    ...  ${POWER_RESTORE_URI}  PowerRestorePolicy  data=${valueDict}
-    ${currentPolicy}=  Read Attribute  ${POWER_RESTORE_URI}  PowerRestorePolicy
+    ...  ${power_restore_uri}  ${policy_attribute}  data=${valueDict}
+    ${currentPolicy}=
+    ...  Read Attribute  ${power_restore_uri}  ${policy_attribute}
     Should Be Equal    ${currentPolicy}   ${policy}
 
 Get System Power Policy