Fix tests for treatments

2 files changed, 57 insertions, 12 deletions

diff --git a/archaeological_finds/tests.py b/archaeological_finds/tests.py
index f609fff47..ae03b2ba4 100644
--- a/archaeological_finds/tests.py
+++ b/archaeological_finds/tests.py
@@ -242,6 +242,13 @@ class TreatmentWizardCreationTest(WizardTest, FindInit, TestCase):
 
         trt_type = models.TreatmentType.objects.get(txt_idx='moving')
         self.form_datas[0].set('basetreatment', 'treatment_type', trt_type.pk)
+        completed, created = models.TreatmentState.objects.get_or_create(
+            txt_idx='completed', defaults={"executed": True, "label": u"Done"}
+        )
+        completed.executed = True
+        completed.save()
+
+        self.form_datas[0].set('basetreatment', 'treatment_state', completed.pk)
 
         self.find, base_find = self.get_default_find(force=True)
         self.form_datas[0].form_datas['selecfind'][
@@ -844,11 +851,18 @@ class TreatmentTest(FindInit, TestCase):
 
         first_find = self.finds[0]
 
-        treatment.save(user=self.get_default_user(), items=self.basket,
-                       treatment_type_list=[treatment_type])
-        treatment.treatment_types.add(treatment_type)
-
+        completed, created = models.TreatmentState.objects.get_or_create(
+            txt_idx='completed', defaults={"executed": True, "label": u"Done"}
+        )
+        completed.executed = True
+        completed.save()
 
+        treatment.treatment_state = completed
+        treatment.save(
+            user=self.get_default_user(), items=self.basket,
+            treatment_type_list=[treatment_type],
+        )
+        treatment.treatment_types.add(treatment_type)
 
         self.assertEqual(items_nb + self.basket.items.count(),
                          models.Find.objects.count(),
@@ -885,8 +899,17 @@ class TreatmentTest(FindInit, TestCase):
         treatment = models.Treatment()
 
         initial_find = self.finds[0]
-        treatment.save(user=self.get_default_user(), items=self.basket,
-                       treatment_type_list=[treatment_type])
+        completed, created = models.TreatmentState.objects.get_or_create(
+            txt_idx='completed', defaults={"executed": True, "label": u"Done"}
+        )
+        completed.executed = True
+        completed.save()
+
+        treatment.treatment_state = completed
+        treatment.save(
+            user=self.get_default_user(), items=self.basket,
+            treatment_type_list=[treatment_type],
+        )
         treatment.treatment_types.add(treatment_type)
 
         self.assertEqual(nb_find, models.Find.objects.count())
@@ -912,8 +935,17 @@ class TreatmentTest(FindInit, TestCase):
         treatment = models.Treatment()
 
         initial_find = self.finds[0]
-        treatment.save(user=self.get_default_user(), items=self.basket,
-                       treatment_type_list=[treatment_type])
+        completed, created = models.TreatmentState.objects.get_or_create(
+            txt_idx='completed', defaults={"executed": True, "label": u"Done"}
+        )
+        completed.executed = True
+        completed.save()
+
+        treatment.treatment_state = completed
+        treatment.save(
+            user=self.get_default_user(), items=self.basket,
+            treatment_type_list=[treatment_type],
+        )
         treatment.treatment_types.add(treatment_type)
 
         nb_b = self.basket.items.count()
@@ -944,8 +976,17 @@ class TreatmentTest(FindInit, TestCase):
         treatment = models.Treatment()
 
         initial_find = self.finds[0]
-        treatment.save(user=self.get_default_user(), items=self.basket,
-                       treatment_type_list=[treatment_type])
+        completed, created = models.TreatmentState.objects.get_or_create(
+            txt_idx='completed', defaults={"executed": True, "label": u"Done"}
+        )
+        completed.executed = True
+        completed.save()
+
+        treatment.treatment_state = completed
+        treatment.save(
+            user=self.get_default_user(), items=self.basket,
+            treatment_type_list=[treatment_type],
+        )
         treatment.treatment_types.add(treatment_type)
 
         nb_b = self.basket.items.count()
diff --git a/ishtar_common/tests.py b/ishtar_common/tests.py
index bd1833594..cf9f599c4 100644
--- a/ishtar_common/tests.py
+++ b/ishtar_common/tests.py
@@ -1235,10 +1235,14 @@ class ShortMenuTest(TestCase):
         self.assertFalse(str(tf.cached_label) in response.content)
 
     def _create_treatment(self):
-        from archaeological_finds.models import Treatment
+        from archaeological_finds.models import Treatment, TreatmentState
+        completed, created = TreatmentState.objects.get_or_create(
+            txt_idx='completed', defaults={"executed": True, "label": u"Done"}
+        )
         return Treatment.objects.create(
             label="My treatment",
-            year=2052
+            year=2052,
+            treatment_state=completed
         )
 
     def test_treatment(self):