Index: suite_leahi/tst_end_of_treatment/test.py
===================================================================
diff -u -ra41aef5f49d9aed41b1d846a4c388459a7c5eed1 -r535c5a49b6d814e8531fa43fc816f0658d1bafef
--- suite_leahi/tst_end_of_treatment/test.py	(.../test.py)	(revision a41aef5f49d9aed41b1d846a4c388459a7c5eed1)
+++ suite_leahi/tst_end_of_treatment/test.py	(.../test.py)	(revision 535c5a49b6d814e8531fa43fc816f0658d1bafef)
@@ -163,15 +163,15 @@
     mouseClick(waitForObject(names.o_createTreatmentRect_TouchRect))
     td.td_operation_mode(TDOpModes.MODE_PRET.value)
     utility.verify_create_treatment_parameters()
-    change_treatmentstates(TDTreatmentStates.TREATMENT_DIALYSIS_STATE.value)
-    change_treatmentstates(TDTreatmentStates.TREATMENT_END_STATE.value)
+    change_treatmentstates(TDTreatmentStates.DIALYSIS_STATE.value)
+    change_treatmentstates(TDTreatmentStates.END_STATE.value)
     verify_stay_treatment_state()
 
     test.startSection("Check the UI behavior when user in end treatment state")
     mouseClick(waitForObject(names.o_treatmentCompleteButton_TreatmentCompleteButton))
     mouseClick(waitForObject(names.o_confirmTouch_ConfirmButton))
     test.verify(waitFor("'endTreatmentState == 6'", 5000), "Testing UI -> TD message end state")
-    change_treatmentstates(TDTreatmentStates.TREATMENT_RINSEBACK_STATE.value)
+    change_treatmentstates(TDTreatmentStates.RINSEBACK_STATE.value)
     test.endSection()