Merge e05b176 into 278afd9

geotrek/infrastructure/management/commands/loadinfrastructure.py

@@ -6,6 +6,7 @@
 from django.core.management.base import BaseCommand, CommandError
 from django.db import transaction
 
+from geotrek.authent.models import default_structure
 from geotrek.authent.models import Structure
 from geotrek.core.helpers import TopologyHelper
 from geotrek.infrastructure.models import Signage, InfrastructureType, InfrastructureCondition, Infrastructure
@@ -36,6 +37,7 @@ def add_arguments(self, parser):
         parser.add_argument('--structure-default', action='store', dest='structure_default', help='Base url')
         parser.add_argument('--description-default', action='store', dest='description_default', default="",
                             help='Base url')
+        parser.add_argument('--eid-field', action='store', dest='eid_field', help='External ID field')
         parser.add_argument('--year-default', action='store', dest='year_default', help='Base url')
 
     def handle(self, *args, **options):
@@ -63,8 +65,10 @@ def handle(self, *args, **options):
         field_structure_type = options.get('structure_field')
         field_description = options.get('description_field')
         field_implantation_year = options.get('year_field')
+        field_eid = options.get('eid_field')
 
         sid = transaction.savepoint()
+        structure_default = options.get('structure_default')
 
         try:
             for layer in data_source:
@@ -85,33 +89,47 @@ def handle(self, *args, **options):
                     self.stdout.write(self.style.ERROR(
                         u"Set it with --name-field, or set a default value with --name-default"))
                     break
-                if (field_condition_type and field_condition_type not in available_fields)\
-                        or (not field_condition_type and not options.get('condition_default')):
+                if field_condition_type and field_condition_type not in available_fields:
                     self.stdout.write(self.style.ERROR(
                         "Field '{}' not found in data source.".format(field_condition_type)))
                     self.stdout.write(self.style.ERROR(
-                        u"Set it with --condition-field, or set a default value with --condition-default"))
+                        u"Change your --condition-field option"))
                     break
-                if (field_structure_type and field_structure_type not in available_fields)\
-                        or (not field_structure_type and not options.get('structure_default')):
+                if field_structure_type and field_structure_type not in available_fields:
                     self.stdout.write(self.style.ERROR(
                         "Field '{}' not found in data source.".format(field_structure_type)))
                     self.stdout.write(self.style.ERROR(
-                        u"Set it with --structure-field, or set a default value with --structure-default"))
+                        u"Change your --structure-field option"))
                     break
-                if (field_description and field_description not in available_fields)\
-                        or (not field_condition_type and not options.get('description_default')):
+                elif not field_structure_type and not structure_default:
+                    structure = default_structure()
+                else:
+                    try:
+                        structure = Structure.objects.get(name=structure_default)
+                        if verbosity > 0:
+                            self.stdout.write(u"Infrastructures will be linked to {}".format(structure))
+                    except Structure.DoesNotExist:
+                        self.stdout.write(u"Infrastructure {} set in options doesn't exist".format(structure_default))
+                        break
+                if field_description and field_description not in available_fields:
                     self.stdout.write(self.style.ERROR(
                         "Field '{}' not found in data source.".format(field_description)))
                     self.stdout.write(self.style.ERROR(
-                        u"Set it with --description-field, or set a default value with --description-default"))
+                        u"Change your --description-field option"))
                     break
-                if (field_implantation_year and field_implantation_year not in available_fields) \
-                        or (not field_condition_type and not options.get('year_default')):
+
+                if field_implantation_year and field_implantation_year not in available_fields:
                     self.stdout.write(
                         self.style.ERROR("Field '{}' not found in data source.".format(field_implantation_year)))
                     self.stdout.write(self.style.ERROR(
-                        "Set it with --implantation-field, or set a default value with --implantation-default"))
+                        u"Change your --year-field option"))
+                    break
+
+                if field_eid and field_eid not in available_fields:
+                    self.stdout.write(
+                        self.style.ERROR("Field '{}' not found in data source.".format(field_eid)))
+                    self.stdout.write(self.style.ERROR(
+                        u"Change your --eid-field option"))
                     break
 
                 for feature in layer:
@@ -122,21 +140,25 @@ def handle(self, *args, **options):
                     type = feature.get(
                         field_infrastructure_type) if field_infrastructure_type in available_fields else options.get(
                         'type_default')
-                    condition = feature.get(
-                        field_condition_type) if field_condition_type in available_fields else options.get(
-                        'condition_default')
-                    structure = feature.get(
-                        field_structure_type) if field_structure_type in available_fields else options.get(
-                        'structure_default')
-                    description = feature.get(field_description) if field_description in available_fields else options.get(
-                        'description_default')
-                    year = int(feature.get(
-                        field_implantation_year)) if field_implantation_year in available_fields else options.get(
-                        'year_default')
-
+                    if field_condition_type in available_fields:
+                        condition = feature.get(field_condition_type)
+                    else:
+                        condition = options.get('condition_default')
+                    structure = Structure.objects.get(feature.get(field_structure_type)) \
+                        if field_structure_type in available_fields else structure
+                    if field_description in available_fields:
+                        description = feature.get(field_description)
+                    else:
+                        description = options.get('description_default')
+                    if field_implantation_year in available_fields:
+                        year = int(feature.get(field_implantation_year))
+                    else:
+                        year = options.get('year_default')
+                    eid = feature.get(field_eid) if field_eid in available_fields else None
                     model = 'S' if options.get('signage') else 'B'
 
-                    self.create_infrastructure(feature_geom, name, type, condition, structure, description, year, model, verbosity)
+                    self.create_infrastructure(feature_geom, name, type, condition, structure, description, year,
+                                               model, verbosity, eid)
 
             transaction.savepoint_commit(sid)
             if verbosity >= 2:
@@ -147,32 +169,41 @@ def handle(self, *args, **options):
             transaction.savepoint_rollback(sid)
             raise
 
-    def create_infrastructure(self, geometry, name, type, condition, structure, description, year, model, verbosity):
+    def create_infrastructure(self, geometry, name, type,
+                              condition, structure, description, year, model, verbosity, eid):
 
         infra_type, created = InfrastructureType.objects.get_or_create(label=type, type=model)
 
         if created and verbosity:
             self.stdout.write(u"- InfrastructureType '{}' created".format(infra_type))
 
-        condition_type, created = InfrastructureCondition.objects.get_or_create(label=condition)
+        if condition:
+            condition_type, created = InfrastructureCondition.objects.get_or_create(label=condition)
+            if created and verbosity:
+                self.stdout.write(u"- Condition Type '{}' created".format(condition_type))
+        else:
+            condition_type = None
 
-        if created and verbosity:
-            self.stdout.write(u"- Condition Type '{}' created".format(condition_type))
-
-        structure, created = Structure.objects.get_or_create(name=structure)
-
-        if created and verbosity:
-            self.stdout.write(u"- Structure '{}' created".format(structure))
         with transaction.atomic():
             Model = Signage if model == 'S' else Infrastructure
-            infra = Model.objects.create(
-                type=infra_type,
-                name=name,
-                condition=condition_type,
-                structure=structure,
-                description=description,
-                implantation_year=year
-            )
+            fields_without_eid = {
+                'type': infra_type,
+                'name': name,
+                'condition': condition_type,
+                'structure': structure,
+                'description': description,
+                'implantation_year': year
+            }
+            if eid:
+                infra, created = Model.objects.update_or_create(
+                    eid=eid,
+                    defaults=fields_without_eid
+                )
+                if verbosity > 0 and not created:
+                    self.stdout.write(u"Update : %s with eid %s" % (name, eid))
+            else:
+                infra = Model.objects.create(**fields_without_eid)
+
         serialized = '{"lng": %s, "lat": %s}' % (geometry.x, geometry.y)
         topology = TopologyHelper.deserialize(serialized)
         infra.mutate(topology)

geotrek/infrastructure/tests/data/line.geojson

@@ -0,0 +1,22 @@
+{
+  "type": "FeatureCollection",
+  "features": [
+    {
+      "type": "Feature",
+      "properties": {},
+      "geometry": {
+        "type": "LineString",
+        "coordinates": [
+          [
+            -3,
+             0
+          ],
+          [
+            -4,
+             0
+          ]
+        ]
+      }
+    }
+  ]
+}

geotrek/infrastructure/tests/test_command_loadinfrastructure.py

@@ -7,13 +7,15 @@
 
 from geotrek.infrastructure.factories import SignageFactory, InfrastructureFactory
 from geotrek.infrastructure.models import Signage, Infrastructure
+from geotrek.authent.factories import StructureFactory
 
 
 class InfrastructureCommandTest(TestCase):
     """
     There are 2 infrastructures in the file signage.shp
     """
     def test_load_signage(self):
+        StructureFactory.create(name='structure')
         filename = os.path.join(os.path.dirname(__file__), 'data', 'signage.shp')
         signage = SignageFactory(name="name", implantation_year=2010)
         call_command('loadinfrastructure', filename, '--signage', type_default='label', name_default='name',
@@ -25,16 +27,39 @@ def test_load_signage(self):
         self.assertEquals(value.count(), 3)
 
     def test_load_infrastructure(self):
+        output = StringIO()
+        structure = StructureFactory.create(name='structure')
         filename = os.path.join(os.path.dirname(__file__), 'data', 'signage.shp')
         building = InfrastructureFactory(name="name", implantation_year=2010)
         call_command('loadinfrastructure', filename, '--infrastructure', type_default='label', name_default='name',
                      condition_default='condition', structure_default='structure',
-                     description_default='description', year_default=2010, verbosity=0)
+                     description_default='description', year_default=2010, verbosity=2, stdout=output)
+        self.assertIn('Infrastructures will be linked to %s' % structure, output.getvalue())
+        self.assertIn('2 objects created.', output.getvalue())
         value = Infrastructure.objects.all()
         self.assertEquals(building.name, value[1].name)
         self.assertEquals(building.implantation_year, value[1].implantation_year)
         self.assertEquals(value.count(), 3)
 
+    def test_load_infrastructure_with_fields(self):
+        output = StringIO()
+        structure = StructureFactory.create(name='structure')
+        filename = os.path.join(os.path.dirname(__file__), 'data', 'signage.shp')
+        InfrastructureFactory(name="name")
+        call_command('loadinfrastructure', filename, '--infrastructure', type_field='label', name_field='name',
+                     condition_field='condition', structure_default='structure',
+                     description_field='descriptio', year_field='year', verbosity=1, stdout=output)
+        self.assertIn('Infrastructures will be linked to %s' % structure, output.getvalue())
+        self.assertIn("InfrastructureType 'type (PNX)' created", output.getvalue())
+        self.assertIn("Condition Type 'condition (PNX)' created", output.getvalue())
+        value = Infrastructure.objects.all()
+        names = [val.name for val in value]
+        years = [val.implantation_year for val in value]
+        self.assertIn('coucou', names)
+        self.assertIn(2010, years)
+        self.assertIn(2012, years)
+        self.assertEquals(value.count(), 3)
+
     def test_no_file_fail(self):
         with self.assertRaises(CommandError) as cm:
             call_command('loadinfrastructure', 'toto.shp')
@@ -47,23 +72,58 @@ def test_load_both_fail(self):
         self.assertEqual(cm.exception.message, "Only one of --signage and --infrastructure required")
 
     def test_missing_defaults(self):
+        StructureFactory.create(name='structure')
         filename = os.path.join(os.path.dirname(__file__), 'data', 'signage.shp')
         output = StringIO()
 
         call_command('loadinfrastructure', filename, '--signage', stdout=output)
         call_command('loadinfrastructure', filename, '--signage', type_default='label', stdout=output)
-        call_command('loadinfrastructure', filename, '--signage', type_default='label', name_default='name',
-                     stdout=output)
-        call_command('loadinfrastructure', filename, '--signage', type_default='label', name_default='name',
-                     condition_default='condition', stdout=output)
-        call_command('loadinfrastructure', filename, '--signage', type_default='label', name_default='name',
-                     condition_default='condition', structure_default='structure', stdout=output)
-        call_command('loadinfrastructure', filename, '--signage', type_default='label', name_default='name',
-                     condition_default='condition', structure_default='structure',
-                     description_default='description', stdout=output)
 
-        elements_to_check = ['type', 'name', 'condition', 'structure', 'description', 'implantation']
-        self.assertEqual(output.getvalue().count("Field 'None' not found in data source."), 6)
+        elements_to_check = ['type', 'name']
+        self.assertEqual(output.getvalue().count("Field 'None' not found in data source."), 2)
         for element in elements_to_check:
             self.assertIn("Set it with --{0}-field, or set a default value with --{0}-default".format(element),
                           output.getvalue())
+
+    def test_wrong_fields_fail(self):
+        StructureFactory.create(name='structure')
+        filename = os.path.join(os.path.dirname(__file__), 'data', 'signage.shp')
+        output = StringIO()
+        call_command('loadinfrastructure', filename, '--signage', type_field='wrong_type_field', stdout=output)
+        call_command('loadinfrastructure', filename, '--signage', type_default='label', name_field='wrong_name_field',
+                     stdout=output)
+        call_command('loadinfrastructure', filename, '--signage', type_default='label', name_field='name',
+                     condition_field='wrong_condition_field', stdout=output)
+        call_command('loadinfrastructure', filename, '--signage', type_default='label', name_field='name',
+                     description_field='wrong_description_field', stdout=output)
+        call_command('loadinfrastructure', filename, '--signage', type_default='label', name_field='name',
+                     year_field='wrong_implantation_year_field', stdout=output)
+        call_command('loadinfrastructure', filename, '--signage', type_default='label', name_field='name',
+                     structure_field='wrong_structure_field', stdout=output)
+        elements_to_check = ['wrong_type_field', 'wrong_name_field', 'wrong_condition_field',
+                             'wrong_description_field', 'wrong_implantation_year_field', 'wrong_structure_field']
+        self.assertEqual(output.getvalue().count("set a default value"), 2)
+        self.assertEqual(output.getvalue().count("Change your"), 4)
+        self.assertEqual(output.getvalue().count("set a default value"), 2)
+        for element in elements_to_check:
+            self.assertIn("Field '{}' not found in data source".format(element),
+                          output.getvalue())
+
+    def test_line_fail_rolling_back(self):
+        StructureFactory.create(name='structure')
+        filename = os.path.join(os.path.dirname(__file__), 'data', 'line.geojson')
+        output = StringIO()
+        with self.assertRaises(IndexError):
+            call_command('loadinfrastructure', filename, '--signage', type_default='label', name_default='name',
+                         stdout=output)
+        self.assertIn('An error occured, rolling back operations.', output.getvalue())
+        self.assertEqual(Infrastructure.objects.count(), 0)
+
+    def test_update_same_eid(self):
+        output = StringIO()
+        filename = os.path.join(os.path.dirname(__file__), 'data', 'signage.shp')
+        SignageFactory(name="name", eid="eid_2")
+        call_command('loadinfrastructure', filename, '--signage', eid_field='eid', type_default='label',
+                     name_default='name', verbosity=2, stdout=output)
+        self.assertIn("Update : name with eid eid1", output.getvalue())
+        self.assertEqual(Signage.objects.count(), 2)