Work on fix for treatment save (refs #3418)

3 files changed, 12 insertions, 5 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: