diff --git a/extended/code_update/update_bmc.robot b/extended/code_update/update_bmc.robot
index 7de2114..205264f 100644
--- a/extended/code_update/update_bmc.robot
+++ b/extended/code_update/update_bmc.robot
@@ -25,6 +25,7 @@
 
 Resource          code_update_utils.robot
 Resource          ../../lib/boot/boot_resource_master.robot
+Resource          ../../lib/state_manager.robot
 
 *** Variables ***
 
@@ -41,10 +42,11 @@
     Check BMC File System Performance
 
 Initiate Code Update BMC
-    [Documentation]    BMC code update process initiation
+    [Documentation]  BMC code update process initiation
+    [Setup]  Set State Interface Version
     [Tags]  Initiate_Code_Update_BMC
 
-    Check If File Exist    ${FILE_PATH}
+    Check If File Exist  ${FILE_PATH}
     System Readiness Test
     ${status}=   Run Keyword and Return Status
     ...   Validate BMC Version   before
@@ -53,7 +55,7 @@
     ...     Pass Execution   Same Driver version installed
 
     Prune Journal Log
-    Initiate Power Off
+    Power Off Request
     Run Keyword And Ignore Error
     ...   Set Policy Setting   RESTORE_LAST_STATE
     Prepare For Update
@@ -62,9 +64,7 @@
     # 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   ${states}
+    Wait For BMC Ready
 
     # TODO: openbmc/openbmc#815
     Sleep  1 min
@@ -85,8 +85,11 @@
     Check If BMC is Up    30 min   10 sec
     Sleep  1 min
     Validate BMC Version
-    Wait Until Keyword Succeeds
-    ...    10 min   10 sec   Verify BMC State   BMC_READY
+
+    # Now that the code update is completed, make sure we use the correct
+    # interface while checking for BMC ready state.
+    Set State Interface Version
+    Wait For BMC Ready
 
 
 Test Basic BMC Performance At Ready State
@@ -97,4 +100,3 @@
     Check BMC Mem Performance
     Check BMC File System Performance
 
-
diff --git a/lib/state_manager.robot b/lib/state_manager.robot
index 97533dc..8177f63 100755
--- a/lib/state_manager.robot
+++ b/lib/state_manager.robot
@@ -8,6 +8,9 @@
 ${BMC_NOT_READY_STATE}       NotReady
 ${QUIET}  ${0}
 
+# "0" indicates that the new "xyz" interface should be used.
+${OBMC_STATES_VERSION}    ${0}
+
 *** Keywords ***
 
 Initiate Host Boot
@@ -132,3 +135,35 @@
     ...    Wait Until Keyword Succeeds
     ...    10 min  10 sec  Is BMC Not Ready
     ...  ELSE  Fail  msg=Invalid BMC state
+
+
+Set State Interface Version
+    [Documentation]  Set version to indicate which interface to use.
+    ${resp}=  Openbmc Get Request  ${CONTROL_URI}chassis0
+    ${status}=  Run Keyword And Return Status
+    ...  Should Be Equal As Strings  ${resp.status_code}  ${HTTP_OK}
+    Run Keyword If  '${status}' == '${True}'
+    ...  Set Global Variable  ${OBMC_STATES_VERSION}  ${1}
+    ...  ELSE
+    ...  Set Global Variable  ${OBMC_STATES_VERSION}  ${0}
+
+
+Power Off Request
+    [Documentation]  Select appropriate poweroff keyword.
+    Run Keyword If  '${OBMC_STATES_VERSION}' == '${1}'
+    ...  Initiate Power Off
+    ...  ELSE
+    ...  Initiate Host PowerOff
+
+
+Wait For BMC Ready
+    [Documentation]  Check BMC state and wait for BMC Ready.
+    @{states}=  Create List  BMC_READY  HOST_POWERED_OFF
+    Run Keyword If  '${OBMC_STATES_VERSION}' == '${1}'
+    ...  Wait Until Keyword Succeeds  10 min  10 sec
+    ...  Verify BMC State  ${states}
+    ...  ELSE
+    ...  Wait Until Keyword Succeeds  10 min  10 sec
+    ...  Is BMC Ready
+
+