Index: tst_main_treatment_ultrafilteration/test.py
===================================================================
diff -u -r653fa5203e3d66003e056f2c3447278bb4a0358c -rbd9b06da4b2fdb60750355489d9004d673ebf9a4
--- tst_main_treatment_ultrafilteration/test.py	(.../test.py)	(revision 653fa5203e3d66003e056f2c3447278bb4a0358c)
+++ tst_main_treatment_ultrafilteration/test.py	(.../test.py)	(revision bd9b06da4b2fdb60750355489d9004d673ebf9a4)
@@ -21,6 +21,7 @@
 from dialin.common.ui_defs import TXStates as uf_states
 from dialin.ui.hd_simulator import HDSimulator
 from configuration import config
+from numpy import diag
 
 TREATMENT_ULTRAFILTERATION_FLOW_PARAMETERS = {
     "OPTION_1" : [100, 1, 1, 1, 1, 1, 1],
@@ -214,7 +215,7 @@
     #navigate to dialysis state
     hd_simulator.cmd_set_treatment_states_data(sub_mode= 2, uf_state= 0, saline_state=0, heparin_state= 0,
                                       rinseback_state= 0, recirculate_state= 0, blood_prime_state= 0,
-                                      treatment_end_state=0, treatment_stop_state= 0) 
+                                      treatment_end_state=0, treatment_stop_state= 0, dialysis_state=0) 
     
     verification_of_uf_from_main_treatment()