diff --git a/tests/test_association.robot b/tests/test_association.robot
index 81c262b..cbec765 100755
--- a/tests/test_association.robot
+++ b/tests/test_association.robot
@@ -5,7 +5,7 @@
 Resource          ../lib/utils.robot
 Resource          ../lib/connection_client.robot
 Resource          ../lib/openbmc_ffdc.robot
-Resource          ../lib/boot/boot_resource_master.robot
+Resource          ../lib/state_manager.robot
 
 Library           Collections
 
@@ -226,15 +226,12 @@
     ${pre_reboot_association_content} =
     ...   Read Attribute   ${association_uri}    endpoints
 
-    Initiate Power Off
-    Check Power Off States
+    Initiate Host PowerOff
 
     ${output}=      Execute Command    /sbin/reboot
     Check If BMC is Up   5 min    10 sec
 
-    @{states}=   Create List   BMC_READY   HOST_POWERED_OFF
-    Wait Until Keyword Succeeds
-    ...    10 min   10 sec   Verify BMC State   ${states}
+    Wait Until Keyword Succeeds  10 min  10 sec  Is BMC Ready
 
     ${post_reboot_association_content} =
     ...   Read Attribute    ${association_uri}    endpoints