Work on fix for treatment save (refs #3418)

1 files changed, 9 insertions, 4 deletions

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 + [