Tests: test container change in treatments

1 files changed, 136 insertions, 17 deletions

diff --git a/archaeological_finds/tests.py b/archaeological_finds/tests.py
index 9efd6ac79..a56c997eb 100644
--- a/archaeological_finds/tests.py
+++ b/archaeological_finds/tests.py
@@ -214,14 +214,15 @@ class TreatmentWizardCreationTest(WizardTest, FindInit, TestCase):
     steps = views.treatment_wizard_steps
     form_datas = [
         FormData(
-            'Move treament',
+            'Move treatment (planned)',
             form_datas={
                 'file': {},
                 'basetreatment': {
                     'treatment_type': None,
                     'person': 1,  # doer
                     'location': 1,  # associated warehouse
-                    'year': 2016
+                    'year': 2016,
+                    'container': None,
                 },
                 'selecfind': {
                     'pk': 1,
@@ -230,7 +231,45 @@ class TreatmentWizardCreationTest(WizardTest, FindInit, TestCase):
             },
             ignored=('resultfind-treatment_creation',
                      'selecbasket-treatment_creation',
-                     'resultfinds-treatment_creation'))
+                     'resultfinds-treatment_creation')),
+        FormData(
+            'Move treatment (done)',
+            form_datas={
+                'file': {},
+                'basetreatment': {
+                    'treatment_type': None,
+                    'person': 1,  # doer
+                    'location': 1,  # associated warehouse
+                    'year': 2016,
+                    'container': None,
+                },
+                'selecfind': {
+                    'pk': 1,
+                    'resulting_pk': 1
+                }
+            },
+            ignored=('resultfind-treatment_creation',
+                     'selecbasket-treatment_creation',
+                     'resultfinds-treatment_creation')),
+        FormData(
+            'Loan treatment (done)',
+            form_datas={
+                'file': {},
+                'basetreatment': {
+                    'treatment_type': None,
+                    'person': 1,  # doer
+                    'location': 1,  # associated warehouse
+                    'year': 2016,
+                    'container': None,
+                },
+                'selecfind': {
+                    'pk': 1,
+                    'resulting_pk': 1
+                }
+            },
+            ignored=('resultfind-treatment_creation',
+                     'selecbasket-treatment_creation',
+                     'resultfinds-treatment_creation')),
     ]
 
     def pre_wizard(self):
@@ -240,38 +279,118 @@ class TreatmentWizardCreationTest(WizardTest, FindInit, TestCase):
                 name="default", warehouse_type=WarehouseType.objects.all()[0])
             warehouse.id = 1
             warehouse.save()
+        else:
+            warehouse = q.all()[0]
         q = Person.objects.filter(pk=1)
         if not q.count():
             person = Person.objects.create(name="default")
             person.id = 1
             person.save()
 
-        trt_type = models.TreatmentType.objects.get(txt_idx='moving')
-        self.form_datas[0].set('basetreatment', 'treatment_type', trt_type.pk)
+        self.container_ref = Container.objects.create(
+            location=warehouse,
+            responsible=warehouse,
+            container_type=ContainerType.objects.all()[0]
+        )
+
+        self.container1 = Container.objects.create(
+            location=warehouse,
+            responsible=warehouse,
+            container_type=ContainerType.objects.all()[0]
+        )
+
+        self.container2 = Container.objects.create(
+            location=warehouse,
+            responsible=warehouse,
+            container_type=ContainerType.objects.all()[0]
+        )
+
+        self.find, base_find = self.get_default_find(force=True)
+        self.find.container_ref = self.container_ref
+        self.find.container = self.container_ref
+        self.find.save()
+        for idx in (0, 1):
+            self.form_datas[idx].form_datas['selecfind']["pk"] = self.find.pk
+            self.form_datas[idx].form_datas['selecfind']["resulting_pk"] = \
+                self.find.pk
+            self.form_datas[idx].set('basetreatment', 'container',
+                                     self.container1.pk)
+
+        self.find_2, base_find = self.get_default_find(force=True)
+        self.find_2.container_ref = self.container_ref
+        self.find_2.container = self.container_ref
+        self.find_2.save()
+        self.form_datas[2].form_datas['selecfind']["pk"] = self.find_2.pk
+        self.form_datas[2].form_datas['selecfind']["resulting_pk"] = \
+            self.find_2.pk
+        self.form_datas[2].set('basetreatment', 'container', self.container2.pk)
+
+        moving = models.TreatmentType.objects.get(txt_idx='moving')
+        moving.change_reference_location = True
+        moving.change_current_location = True
+        moving.save()
+        self.form_datas[0].set('basetreatment', 'treatment_type',
+                               moving.pk)
+        self.form_datas[1].set('basetreatment', 'treatment_type',
+                               moving.pk)
+        loan = models.TreatmentType.objects.get(txt_idx='loan')
+        loan.change_reference_location = False
+        loan.change_current_location = True
+        loan.save()
+        self.form_datas[2].set('basetreatment', 'treatment_type',
+                               loan.pk)
+
+        planned, __ = models.TreatmentState.objects.get_or_create(
+            txt_idx='planned', defaults={"executed": False, "label": u"Planned"}
+        )
+        planned.executed = False
+        planned.save()
+        self.form_datas[0].set('basetreatment', 'treatment_state', planned.pk)
+
         completed, created = models.TreatmentState.objects.get_or_create(
             txt_idx='completed', defaults={"executed": True, "label": u"Done"}
         )
         completed.executed = True
         completed.save()
+        self.form_datas[1].set('basetreatment', 'treatment_state', completed.pk)
+        self.form_datas[2].set('basetreatment', 'treatment_state', completed.pk)
 
-        self.form_datas[0].set('basetreatment', 'treatment_state', completed.pk)
-
-        self.find, base_find = self.get_default_find(force=True)
-        self.form_datas[0].form_datas['selecfind'][
-            'pk'] = self.find.pk
-        self.form_datas[0].form_datas['selecfind'][
-            'resulting_pk'] = self.find.pk
         self.treatment_number = models.Treatment.objects.count()
+
+        def post_first_wizard(test_object, final_step_response):
+            # treatment planned - no changes
+            find = models.Find.objects.get(pk=test_object.find.pk)
+            test_object.assertEqual(find.container, test_object.container_ref)
+            test_object.assertEqual(find.container_ref,
+                                    test_object.container_ref)
+        self.form_datas[0].extra_tests = [post_first_wizard]
+
+        def post_second_wizard(test_object, final_step_response):
+            # moving done
+            find = models.Find.objects.get(pk=test_object.find.pk)
+            test_object.assertEqual(find.container, test_object.container1)
+            test_object.assertEqual(find.container_ref, test_object.container1)
+        self.form_datas[1].extra_tests = [post_second_wizard]
+
+        def post_third_wizard(test_object, final_step_response):
+            # loan done
+            find = models.Find.objects.get(pk=test_object.find_2.pk)
+            test_object.assertEqual(find.container_ref,
+                                    test_object.container_ref)
+            test_object.assertEqual(find.container, test_object.container2)
+        self.form_datas[2].extra_tests = [post_third_wizard]
+
         super(TreatmentWizardCreationTest, 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.treatment_number + 3)
         self.find = models.Find.objects.get(pk=self.find.pk)
-        self.assertEqual(models.Find.objects.filter(
-            treatments=treat).count(), 1)
+        self.assertEqual(self.find.treatments.count(), 2)
+        self.find_2 = models.Find.objects.get(pk=self.find_2.pk)
+        self.assertEqual(self.find_2.treatments.count(), 1)
         # TODO: test treatment with new find creation
+        # treat = models.Treatment.objects.order_by('-pk').all()[0]
         # self.assertEqual(self.find.downstream_treatment,
         #                 treat)
 
@@ -429,7 +548,7 @@ class FindTest(FindInit, TestCase):
         profile = get_current_profile()
         profile.find_index = u"O"
         profile.save()
-        profile = get_current_profile(force=True)
+        get_current_profile(force=True)
 
         op1 = self.create_operation()[-1]
         op2 = self.create_operation()[-1]