Work on fix for treatment save (refs #3418)

4 files changed, 19 insertions, 7 deletions

diff --git a/archaeological_finds/models_treatments.py b/archaeological_finds/models_treatments.py
index 90badcd35..dacfbfdfc 100644
--- a/archaeological_finds/models_treatments.py
+++ b/archaeological_finds/models_treatments.py
@@ -90,6 +90,8 @@ class Treatment(BaseHistorizedItem, ImageModel, OwnPerms, ShortMenuItem):
         "upstream_cached_label": _(u"Upstream find"),
     }
     IMAGE_PREFIX = 'treatment'
+    # extra keys than can be passed to save method
+    EXTRA_SAVED_KEYS = ('items', 'user')
     SLUG = 'treatment'
     label = models.CharField(_(u"Label"), blank=True, null=True,
                              max_length=200)
diff --git a/archaeological_finds/tests.py b/archaeological_finds/tests.py
index a1410c5d9..be83de77b 100644
--- a/archaeological_finds/tests.py
+++ b/archaeological_finds/tests.py
@@ -26,7 +26,6 @@ from ishtar_common.models import ImporterType, IshtarUser, ImporterColumn,\
     FormaterType, ImportTarget
 
 from ishtar_common.models import Person
-from archaeological_context_records.models import ContextRecord
 from archaeological_finds import models, views
 from archaeological_warehouse.models import Warehouse, WarehouseType
 
@@ -187,20 +186,26 @@ class ATreatmentWizardCreationTest(WizardTest, FindInit, TestCase):
             person = Person.objects.create(name="default")
             person.id = 1
             person.save()
-        find, base_find = self.get_default_find(force=True)
+        self.find, base_find = self.get_default_find(force=True)
         self.form_datas[0].form_datas['selecfind-treatment_creation'][
-            'pk'] = find.pk
+            'pk'] = self.find.pk
         self.form_datas[0].form_datas['selecfind-treatment_creation'][
-            'resulting_pk'] = find.pk
+            'resulting_pk'] = self.find.pk
         self.treatment_number = models.Treatment.objects.count()
         super(ATreatmentWizardCreationTest, self).pre_wizard()
 
     def post_wizard(self):
         self.assertEqual(models.Treatment.objects.count(),
                          self.treatment_number + 1)
+        treat = models.Treatment.objects.order_by('-pk').all()[0]
+        self.assertEqual(models.Find.objects.filter(
+            upstream_treatment=treat).count(), 1)
+        self.assertEqual(self.find.downstream_treatment,
+                         treat)
 
 
 class ImportFindTest(ImportContextRecordTest):
+    test_operations = False
     test_context_records = False
 
     fixtures = ImportContextRecordTest.fixtures + [
diff --git a/archaeological_finds/wizards.py b/archaeological_finds/wizards.py
index dab4a76d4..96a83e582 100644
--- a/archaeological_finds/wizards.py
+++ b/archaeological_finds/wizards.py
@@ -101,7 +101,7 @@ class TreatmentWizard(Wizard):
 
     def get_extra_model(self, dct, form_list):
         """
-        Remove basket ID to the result dict
+        Get items concerned by the treatment
         """
         dct = super(TreatmentWizard, self).get_extra_model(dct, form_list)
         if 'resulting_pk' in dct:
diff --git a/ishtar_common/wizards.py b/ishtar_common/wizards.py
index ce1333c89..87a126879 100644
--- a/ishtar_common/wizards.py
+++ b/ishtar_common/wizards.py
@@ -571,7 +571,10 @@ class Wizard(NamedUrlWizardView):
             all_field_names = self.get_saved_model()._meta.get_all_field_names()
             for k in dct.copy():
                 if not (k.endswith('_id') and k[:-3] in all_field_names) \
-                        and k not in all_field_names:
+                        and k not in all_field_names and \
+                        (not hasattr(self.get_saved_model(),
+                                     'EXTRA_SAVED_KEYS') or
+                         k not in self.get_saved_model().EXTRA_SAVED_KEYS):
                     dct.pop(k)
             saved_args = self.saved_args.copy()
             for k in saved_args:
@@ -687,7 +690,9 @@ class Wizard(NamedUrlWizardView):
         return return_object and (obj, res) or res
 
     def get_deleted(self, keys):
-        """Get the deleted and non-deleted items in formsets"""
+        """
+        Get the deleted and non-deleted items in formsets
+        """
         not_to_delete, to_delete = set(), set()
         for key in keys:
             items = key.split('-')