Treatment: workon on basket update after treatment

3 files changed, 112 insertions, 12 deletions

diff --git a/archaeological_finds/models.py b/archaeological_finds/models.py
index c8b392310..0a1e59a68 100644
--- a/archaeological_finds/models.py
+++ b/archaeological_finds/models.py
@@ -612,6 +612,35 @@ class Treatment(BaseHistorizedItem, OwnPerms):
             lbl += u" %s %s" % (_(u"by"), unicode(self.person))
         return lbl
 
+    def save(self, *args, **kwargs):
+        items, user, extra_args_for_new = [], None, []
+        if "items" in kwargs:
+            items = kwargs.pop('items')
+        if "user" in kwargs:
+            user = kwargs.pop('user')
+        if "extra_args_for_new" in kwargs:
+            extra_args_for_new = kwargs.pop('extra_args_for_new')
+        is_new = self.pk is None
+        super(Treatment, self).save(*args, **kwargs)
+        if not is_new or not items:
+            return
+        basket = None
+        if hasattr(items, "items"):
+            basket = items
+            items = basket.items.all()
+        for item in items:
+            new = item.duplicate(user)
+            item.downstream_treatment = self
+            item.save()
+            new.upstream_treatment = self
+            for k in extra_args_for_new:
+                setattr(new, k, extra_args_for_new[k])
+            new.save()
+            # update baskets
+            for basket in FindBasket.objects.filter(items__pk=item.pk).all():
+                basket.items.remove(item)
+                basket.items.add(new)
+
 
 class TreatmentSource(Source):
     class Meta:
diff --git a/archaeological_finds/tests.py b/archaeological_finds/tests.py
index ccf821757..f75e9cfa3 100644
--- a/archaeological_finds/tests.py
+++ b/archaeological_finds/tests.py
@@ -19,11 +19,13 @@
 
 from django.conf import settings
 from django.core.files.uploadedfile import SimpleUploadedFile
-from ishtar_common.models import ImporterType
+from django.test import TestCase
+from ishtar_common.models import ImporterType, IshtarUser
 
 from archaeological_finds import models
 
-from archaeological_context_records.tests import ImportContextRecordTest
+from archaeological_context_records.tests import ImportContextRecordTest, \
+    ContextRecordInit
 
 from ishtar_common import forms_common
 
@@ -67,3 +69,78 @@ class ImportFindTest(ImportContextRecordTest):
         self.assertTrue(current_nb == (old_nb + 4))
         self.assertEqual(
             models.Find.objects.filter(material_types__pk=ceram).count(), 4)
+
+
+class FindInit(ContextRecordInit):
+    test_context_records = False
+
+    def create_finds(self, user=None, data_base={}, data={}):
+        if not getattr(self, 'finds', None):
+            self.finds = []
+        if not getattr(self, 'base_finds', None):
+            self.base_finds = []
+
+        default = {'label': "Base find"}
+        if not data_base.get('history_modifier'):
+            data_base['history_modifier'] = self.get_default_user()
+        if not data_base.get('context_record'):
+            data_base['context_record'] = self.get_default_context_record()
+
+        default.update(data_base)
+        base_find = models.BaseFind.objects.create(**default)
+        self.base_finds.append(base_find)
+
+        data["history_modifier"] = data_base["history_modifier"]
+        find = models.Find.objects.create(**data)
+        find.base_finds.add(base_find)
+        self.finds.append(find)
+        return self.finds, self.base_finds
+
+    def get_default_find(self):
+        return self.create_finds()[0]
+
+
+class PackagingTest(TestCase, FindInit):
+    fixtures = [settings.ROOT_PATH +
+                '../fixtures/initial_data.json',
+                settings.ROOT_PATH +
+                '../ishtar_common/fixtures/initial_data.json',
+                settings.ROOT_PATH +
+                '../archaeological_files/fixtures/initial_data.json',
+                settings.ROOT_PATH +
+                '../archaeological_operations/fixtures/initial_data-fr.json',
+                settings.ROOT_PATH +
+                '../archaeological_finds/fixtures/initial_data-fr.json',
+                ]
+    model = models.Find
+
+    def setUp(self):
+        self.create_finds({"label": u"Find 1"})
+        self.create_finds({"label": u"Find 2"})
+        self.basket = models.FindBasket.objects.create(
+            label="My basket", user=IshtarUser.objects.get(
+                pk=self.get_default_user().pk))
+        self.other_basket = models.FindBasket.objects.create(
+            label="My other basket", user=IshtarUser.objects.get(
+                pk=self.get_default_user().pk))
+        for find in self.finds:
+            self.basket.items.add(find)
+            self.other_basket.items.add(find)
+
+    def testPackaging(self):
+        treatment_type = models.TreatmentType.objects.get(txt_idx='packaging')
+        treatment = models.Treatment(treatment_type=treatment_type)
+        items_nb = models.Find.objects.count()
+        treatment.save(user=self.get_default_user(), items=self.basket)
+        self.assertEqual(items_nb + self.basket.items.count(),
+                         models.Find.objects.count(),
+                         msg="Packaging doesn't generate enough new finds")
+        # new version of the find is in the basket
+        for item in self.basket.items.all():
+            self.assertNotIn(
+                item, self.finds,
+                msg="Original basket have not been upgraded after packaging")
+        for item in self.other_basket.items.all():
+            self.assertNotIn(
+                item, self.finds,
+                msg="Other basket have not been upgraded after packaging")
diff --git a/archaeological_warehouse/wizards.py b/archaeological_warehouse/wizards.py
index 5f2b4346d..409f7b28f 100644
--- a/archaeological_warehouse/wizards.py
+++ b/archaeological_warehouse/wizards.py
@@ -33,22 +33,16 @@ class PackagingWizard(TreatmentWizard):
         kwargs['user'] = self.request.user
         return kwargs
 
-
     def save_model(self, dct, m2m, whole_associated_models, form_list,
                    return_object):
         dct = self.get_extra_model(dct, form_list)
         obj = self.get_current_saved_object()
         dct['location'] = dct['container'].location
-        basket = dct.pop('basket')
+        items = dct.pop('basket')
         treatment = Treatment(**dct)
-        treatment.save()
-        for item in basket.items.all():
-            new = item.duplicate(self.request.user)
-            item.downstream_treatment = treatment
-            item.save()
-            new.upstream_treatment = treatment
-            new.container = dct['container']
-            new.save()
+        extra_args_for_new = {"container": dct['container']}
+        treatment.save(items=items, user=self.request.user,
+                       extra_args_for_new=extra_args_for_new)
         res = render_to_response('ishtar/wizard/wizard_done.html', {},
                                  context_instance=RequestContext(self.request))
         return return_object and (obj, res) or res