Fix division importer

2 files changed, 25 insertions, 15 deletions

diff --git a/archaeological_warehouse/models.py b/archaeological_warehouse/models.py
index 71751fcbe..3b319881d 100644
--- a/archaeological_warehouse/models.py
+++ b/archaeological_warehouse/models.py
@@ -386,11 +386,11 @@ class Warehouse(Address, DocumentItem, GeoItem, CompleteIdentifierItem,
         return current_container_type, previous_container_types
 
     @post_importer_action
-    def add_localisations(self, __, value):
-        self._add_localisations(value)
+    def add_localisations(self, context, value):
+        self._add_localisations(context, value)
     add_localisations.post_save = True
 
-    def _add_localisations(self, value, return_errors=False):
+    def _add_localisations(self, context, value, return_errors=False):
         """
         Add localisations for this warehouse.
         Get the default localisation types and set each reference from the
@@ -406,6 +406,10 @@ class Warehouse(Address, DocumentItem, GeoItem, CompleteIdentifierItem,
                 return None, _("No value")
             return
 
+        import_object = None
+        if context and "import_object" in context:
+            import_object = context["import_object"]
+
         TMP_SEMI_COLON = "|#|#|"
         value = value.replace("\\;", TMP_SEMI_COLON)  # manage ";" used by a ref
 
@@ -413,23 +417,27 @@ class Warehouse(Address, DocumentItem, GeoItem, CompleteIdentifierItem,
 
         divisions = list(WarehouseDivisionLink.objects.filter(
             warehouse=self).order_by('order'))
-        if len(values) > len(divisions):
-            if return_errors:
-                return str(_("{} values for only {} default divisions set for "
-                             "warehouse {}")).format(
-                    len(values), len(divisions), self.name)
-            return
 
         parent = None
         for idx, value in enumerate(values):
+            if idx >= len(divisions):
+                if return_errors:
+                    return str(
+                        _("{} values for only {} default divisions set for "
+                          "warehouse {}")).format(
+                        len(values), len(divisions), self.name)
+                return
             value = value.replace(TMP_SEMI_COLON, ";").strip()
             if not value or value == "-":
                 continue
-            parent, __ = Container.objects.get_or_create(
+            parent, created = Container.objects.get_or_create(
                 location=self,
                 reference=value,
                 container_type_id=divisions[idx].container_type_id,
                 parent=parent)
+            if created and import_object:
+                parent.imports.add(import_object)
+
 
     @property
     def short_label(self):
diff --git a/archaeological_warehouse/tests.py b/archaeological_warehouse/tests.py
index 265a4a7ef..25937ed6a 100644
--- a/archaeological_warehouse/tests.py
+++ b/archaeological_warehouse/tests.py
@@ -500,7 +500,7 @@ class WarehouseTest(TestCase):
         container_nb = models.Container.objects.count()
         base_value = "A;42;allée 3\;2"
         error = self.warehouse._add_localisations(
-            base_value, return_errors=True)
+            None, base_value, return_errors=True)
         self.assertTrue(error)  # no division set
 
         for idx, ct in enumerate(self.container_types):
@@ -508,11 +508,11 @@ class WarehouseTest(TestCase):
                 container_type=ct, order=idx * 10, warehouse=self.warehouse)
         too_many_value = "A;42;allée 3\;2;5;42;3"
         error = self.warehouse._add_localisations(
-            too_many_value, return_errors=True)
+            None, too_many_value, return_errors=True)
         self.assertTrue(error)
 
         error = self.warehouse._add_localisations(
-            base_value, return_errors=True)
+            None, base_value, return_errors=True)
         self.assertIsNone(error)
         parent = None
         for idx, reference in enumerate(("A", "42", "allée 3;2")):
@@ -529,7 +529,8 @@ class WarehouseTest(TestCase):
         self.assertEqual(container_nb + 3, new_container_nb)
 
         value = "A;42;allée 4"
-        error = self.warehouse._add_localisations(value, return_errors=True)
+        error = self.warehouse._add_localisations(
+            None, value, return_errors=True)
         self.assertIsNone(error)
         # only create a new container
         self.assertEqual(new_container_nb + 1, models.Container.objects.count())
@@ -544,7 +545,8 @@ class WarehouseTest(TestCase):
 
         # test with an empty localisation
         value = "A;42;;35"
-        error = self.warehouse._add_localisations(value, return_errors=True)
+        error = self.warehouse._add_localisations(
+            None, value, return_errors=True)
         self.assertIsNone(error)
         q = models.Container.objects.filter(
             parent__reference="42",