fix: avoid exception when country code is not found (#1074)

Author: davidgamez

api/src/scripts/populate_db_gtfs.py

@@ -1,12 +1,14 @@
 import os
 import traceback
-from typing import TYPE_CHECKING
 from datetime import datetime
+from typing import TYPE_CHECKING
 
 import pycountry
 import pytz
 from sqlalchemy import text
 
+from scripts.load_dataset_on_create import publish_all
+from scripts.populate_db import DatabasePopulateHelper, set_up_configs
 from shared.database.database import generate_unique_id, configure_polymorphic_mappers
 from shared.database_gen.sqlacodegen_models import (
     Entitytype,
@@ -16,8 +18,6 @@
     Redirectingid,
     t_feedsearch,
 )
-from scripts.populate_db import DatabasePopulateHelper, set_up_configs
-from scripts.load_dataset_on_create import publish_all
 from utils.data_utils import set_up_defaults
 
 if TYPE_CHECKING:
@@ -61,9 +61,11 @@ def get_stable_id(self, row):
         return f'mdb-{self.get_safe_value(row, "mdb_source_id", "")}'
 
     def get_country(self, country_code):
+        country = None
         if country_code:
-            return pycountry.countries.get(alpha_2=country_code).name
-        return None
+            country = pycountry.countries.get(alpha_2=country_code)
+            country = country.name if country else None
+        return country
 
     def populate_location(self, session, feed, row, stable_id):
         """
@@ -88,7 +90,10 @@ def populate_location(self, session, feed, row, stable_id):
                 if location
                 else Location(
                     id=location_id,
-                    country_code=country_code,
+                    # Country code should be short.
+                    # If too long it might be an error
+                    # (like it could be the country name instead of code).
+                    country_code=country_code if country_code and len(country_code) <= 3 else None,
                     subdivision_name=subdivision_name,
                     municipality=municipality,
                     country=country,

functions-python/reverse_geolocation/src/location_group_utils.py

@@ -75,7 +75,8 @@ def __init__(self, location_group: Osmlocationgroup, stops_count: int):
 
     def country(self) -> str:
         """Returns the country name of the LocationGroup."""
-        return pycountry.countries.get(alpha_2=self.iso_3166_1_code).name
+        country = pycountry.countries.get(alpha_2=self.iso_3166_1_code)
+        return country.name if country else None
 
     def location_id(self) -> str:
         """Returns the location ID of the LocationGroup."""

functions-python/reverse_geolocation/tests/test_location_group_utils.py

@@ -2,12 +2,13 @@
 import unittest
 from unittest.mock import patch, MagicMock
 
+import pytest
 from faker import Faker
 from geoalchemy2 import WKTElement
 
+from location_group_utils import GeopolygonAggregate
 from shared.database_gen.sqlacodegen_models import Geopolygon, Osmlocationgroup
 
-
 faker = Faker()
 
 
@@ -91,3 +92,26 @@ def test_geopolygon_aggregate(self):
         geopolygon_aggregate_2 = GeopolygonAggregate(location_group, 1)
         geopolygon_aggregate.merge(geopolygon_aggregate_2)
         self.assertEqual(geopolygon_aggregate.stop_count, 2)
+
+
+@pytest.mark.parametrize(
+    "values",
+    [{"value": "CA", "expected": "Canada"}, {"value": "Canada", "expected": None}],
+)
+def test_location_country(values):
+    """
+    Test the country function with cases with valid and invalid ISO 3166_1 code
+    """
+    geopolygon = Geopolygon(
+        osm_id=1,
+        admin_level=2,
+        name=values.get("value"),
+        iso_3166_1_code=values.get("value"),
+        geometry=WKTElement("POINT(-73.5673 45.5017)", srid=4326),
+    )
+    location_group = Osmlocationgroup(
+        group_id="1.1.1", group_name="Canada, Ontario", osms=[geopolygon]
+    )
+    geopolygon_aggregate = GeopolygonAggregate(location_group, 1)
+
+    assert geopolygon_aggregate.country() == values.get("expected")