Fix packaging creation and treatment tests

2 files changed, 26 insertions, 8 deletions

diff --git a/archaeological_finds/models_treatments.py b/archaeological_finds/models_treatments.py
index 48da7c8ec..bc5eace58 100644
--- a/archaeological_finds/models_treatments.py
+++ b/archaeological_finds/models_treatments.py
@@ -470,7 +470,11 @@ class Treatment(DashboardFormItem, ValueGetter, BaseHistorizedItem,
         # manage containers
         if not self.container:
             return
-        for find in Find.objects.filter(upstream_treatment=self).all():
+        if create_new_find:
+            q = Find.objects.filter(upstream_treatment=self)
+        else:
+            q = Find.objects.filter(treatments=self)
+        for find in q.all():
             if find.container != self.container:
                 find.container = self.container
                 if find.pk in updated:
diff --git a/archaeological_finds/tests.py b/archaeological_finds/tests.py
index a55e075b7..b5ec5de83 100644
--- a/archaeological_finds/tests.py
+++ b/archaeological_finds/tests.py
@@ -258,9 +258,10 @@ class TreatmentWizardCreationTest(WizardTest, FindInit, TestCase):
         treat = models.Treatment.objects.order_by('-pk').all()[0]
         self.find = models.Find.objects.get(pk=self.find.pk)
         self.assertEqual(models.Find.objects.filter(
-            upstream_treatment=treat).count(), 1)
-        self.assertEqual(self.find.downstream_treatment,
-                         treat)
+            treatments=treat).count(), 1)
+        # TODO: test treatment with new find creation
+        # self.assertEqual(self.find.downstream_treatment,
+        #                 treat)
 
 
 class ImportFindTest(ImportTest, TestCase):
@@ -832,18 +833,26 @@ class PackagingTest(FindInit, TestCase):
             self.basket.items.add(find)
             self.other_basket.items.add(find)
 
-    def testPackaging(self):
+    def test_packaging_with_new_find_creation(self):
         treatment_type = models.TreatmentType.objects.get(txt_idx='packaging')
+        # make packaging a treatment with a new version of the find created
+        treatment_type.create_new_find = True
+        treatment_type.save()
+
         treatment = models.Treatment()
         items_nb = models.Find.objects.count()
 
         first_find = self.finds[0]
 
-        treatment.save(user=self.get_default_user(), items=self.basket)
+        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(),
                          msg="Packaging doesn't generate enough new finds")
-        treatment.treatment_types.add(treatment_type)
 
         resulting_find = models.Find.objects.get(
             upstream_treatment__upstream=first_find,
@@ -869,10 +878,15 @@ class PackagingTest(FindInit, TestCase):
     def test_delete(self):
         # manage treatment deletion
         treatment_type = models.TreatmentType.objects.get(txt_idx='packaging')
+        # make packaging a treatment with a new version of the find created
+        treatment_type.create_new_find = True
+        treatment_type.save()
+
         treatment = models.Treatment()
 
         initial_find = self.finds[0]
-        treatment.save(user=self.get_default_user(), items=self.basket)
+        treatment.save(user=self.get_default_user(), items=self.basket,
+                       treatment_type_list=[treatment_type])
         treatment.treatment_types.add(treatment_type)
 
         resulting_find = models.Find.objects.get(