Context record wizard: fix related context record association on creation

3 files changed, 75 insertions, 4 deletions

diff --git a/archaeological_context_records/tests.py b/archaeological_context_records/tests.py
index df05f4898..07360f73a 100644
--- a/archaeological_context_records/tests.py
+++ b/archaeological_context_records/tests.py
@@ -27,13 +27,16 @@ from django.core.urlresolvers import reverse
 from django.test.client import Client
 
 from ishtar_common.models import IshtarSiteProfile, ImporterModel
-from ishtar_common.tests import create_superuser, TestCase
 
 from archaeological_operations.tests import OperationInitTest, \
     ImportTest, ImportOperationTest
 from archaeological_operations import models as models_ope
 from archaeological_context_records import models
 
+from ishtar_common.tests import WizardTest, WizardTestFormData as FormData, \
+    create_superuser, TestCase
+
+from archaeological_context_records import views
 
 class ImportContextRecordTest(ImportTest, TestCase):
 
@@ -527,3 +530,66 @@ class RecordRelationsTest(ContextRecordInit, TestCase):
         self.assertEqual(models.RecordRelations.objects.filter(
             left_record=cr_2, right_record=cr_1,
             relation_type=rel_type_2).count(), 1)
+
+
+class ContextRecordWizardCreationTest(WizardTest, ContextRecordInit, TestCase):
+    fixtures = ImportOperationTest.fixtures
+    url_name = 'record_creation'
+    wizard_name = 'record_wizard'
+    steps = views.record_creation_steps
+    form_datas = [
+        FormData(
+            "Create a simple context record",
+            form_datas={
+                'selec-record_creation': {},
+                'general-record_creation': {
+                    'label': "First"
+                },
+            },
+            ignored=('datings-record_creation',
+                     'interpretation-record_creation',
+                     'relations-record_creation',)
+        ),
+        FormData(
+            "Create a context record with a relation",
+            form_datas={
+                'selec-record_creation': {},
+                'general-record_creation': {
+                    'label': "Second"
+                },
+                'relations-record_creation': []
+            },
+            ignored=('datings-record_creation',
+                     'interpretation-record_creation',)
+        ),
+    ]
+
+    def pre_wizard(self):
+        profile, created = IshtarSiteProfile.objects.get_or_create(
+            slug='default', active=True)
+        profile.context_record = True
+        profile.save()
+
+        ope = self.get_default_operation()
+        parcel = self.get_default_parcel()
+        for form_data in self.form_datas:
+            form_data.form_datas['selec-record_creation']['operation_id'] = \
+                ope.pk
+            form_data.form_datas['general-record_creation']['parcel'] = \
+                parcel.pk
+
+        self.related_cr = self.create_context_record(data={'operation': ope})[0]
+        self.form_datas[1].form_datas['relations-record_creation'].append(
+            {'right_record': self.related_cr.pk,
+             'relation_type': models.RelationType.objects.create(
+                 label="Test").pk}
+        )
+
+        self.cr_nb = models.ContextRecord.objects.count()
+        super(ContextRecordWizardCreationTest, self).pre_wizard()
+
+    def post_wizard(self):
+        self.assertEqual(models.ContextRecord.objects.count(),
+                         self.cr_nb + 2)
+        self.assertEqual(self.related_cr.left_relations.count(),
+                         1)
diff --git a/archaeological_context_records/views.py b/archaeological_context_records/views.py
index cf3a13b3b..250fdafc0 100644
--- a/archaeological_context_records/views.py
+++ b/archaeological_context_records/views.py
@@ -92,13 +92,16 @@ record_search_wizard = SearchWizard.as_view([
     label=_(u"Context record search"),
     url_name='record_search',)
 
-record_creation_wizard = RecordWizard.as_view([
+record_creation_steps = [
     ('selec-record_creation', OperationRecordFormSelection),
     ('general-record_creation', RecordFormGeneral),
     ('datings-record_creation', DatingFormSet),
     ('interpretation-record_creation', RecordFormInterpretation),
     ('relations-record_creation', RecordRelationsFormSet),
-    ('final-record_creation', FinalForm)],
+    ('final-record_creation', FinalForm)
+]
+
+record_creation_wizard = RecordWizard.as_view(record_creation_steps,
     label=_(u"New context record"),
     url_name='record_creation',)
 
diff --git a/archaeological_operations/forms.py b/archaeological_operations/forms.py
index 314e184e2..058e637f4 100644
--- a/archaeological_operations/forms.py
+++ b/archaeological_operations/forms.py
@@ -416,7 +416,9 @@ class RecordRelationsForm(ManageOldType, forms.Form):
                 cleaned_data.get('right_record', None)):
             raise forms.ValidationError(
                 _(u"You should select a relation type."))
-        if str(cleaned_data.get('right_record')) == str(self.left_record.pk):
+        if self.left_record and \
+                str(cleaned_data.get('right_record')) == str(
+                    self.left_record.pk):
             raise forms.ValidationError(
                 _(u"An operation cannot be related to herself."))
         return cleaned_data