Fix packaging creation and treatment tests

4 files changed, 30 insertions, 9 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(
diff --git a/archaeological_warehouse/urls.py b/archaeological_warehouse/urls.py
index fa895b7c4..947129c82 100644
--- a/archaeological_warehouse/urls.py
+++ b/archaeological_warehouse/urls.py
@@ -29,7 +29,8 @@ from archaeological_warehouse import models
 # forms
 urlpatterns = [
     url(r'warehouse_packaging/(?P<step>.+)?$',
-        views.warehouse_packaging_wizard, name='warehouse_packaging'),
+        check_rights(['change_find', 'change_own_find'])(
+        views.warehouse_packaging_wizard), name='warehouse_packaging'),
 
     url(r'new-warehouse/(?P<parent_name>.+)?/$',
         views.new_warehouse, name='new-warehouse'),
diff --git a/archaeological_warehouse/wizards.py b/archaeological_warehouse/wizards.py
index a0d46e892..03258f1e2 100644
--- a/archaeological_warehouse/wizards.py
+++ b/archaeological_warehouse/wizards.py
@@ -60,6 +60,8 @@ class PackagingWizard(TreatmentWizard):
                 items = dct.pop('basket')
             else:
                 dct.pop('basket')
+        if 'treatment_type_list' in dct:
+            dct.pop('treatment_type_list')
         treatment = Treatment(**dct)
         extra_args_for_new = {"container": dct['container']}
         treatment.save(items=items, user=self.request.user,