fix tests imports

Author: davidgamez

functions-python/batch_datasets/src/main.py

@@ -29,7 +29,8 @@
 from sqlalchemy.orm import Session
 
 from shared.database_gen.sqlacodegen_models import Gtfsfeed, Gtfsdataset
-from shared.dataset_service.main import BatchExecutionService, BatchExecution
+from shared.dataset_service.dataset_service_commons import BatchExecution
+from shared.dataset_service.main import BatchExecutionService
 from shared.database.database import with_db_session
 from shared.helpers.logger import init_logger
 

functions-python/dataset_service/main.py

@@ -13,19 +13,30 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 #
+import importlib
 import logging
 import uuid
 from dataclasses import asdict
 from typing import Final
 from google.cloud import datastore
 from google.cloud.datastore import Client
 
-from dataset_service_commons import (
-    DatasetTrace,
-    Status,
-    PipelineStage,
-    BatchExecution,
-)
+# This allows the module to be run as a script or imported as a module
+if __package__ is None or __package__ == "":
+    import os
+    import sys
+
+    sys.path.append(os.path.dirname(os.path.abspath(__file__)))
+    import dataset_service_commons
+else:
+    dataset_service_commons = importlib.import_module(
+        ".dataset_service_commons", package=__package__
+    )
+
+Status = dataset_service_commons.Status
+PipelineStage = dataset_service_commons.PipelineStage
+BatchExecution = dataset_service_commons.BatchExecution
+DatasetTrace = dataset_service_commons.DatasetTrace
 
 # This files contains the dataset trace and batch execution models and services.
 # The dataset trace is used to store the trace of a dataset and the batch execution

functions-python/dataset_service/tests/__init__.py

@@ -2,13 +2,8 @@
 from datetime import datetime
 from unittest.mock import patch, MagicMock
 
-from main import (
-    DatasetTrace,
-    DatasetTraceService,
-    Status,
-    BatchExecutionService,
-    BatchExecution,
-)
+from dataset_service_commons import DatasetTrace, Status, BatchExecution
+from main import DatasetTraceService, BatchExecutionService
 
 
 class TestDatasetService(unittest.TestCase):

functions-python/dataset_service/tests/test_dataset_service.py

@@ -28,9 +28,6 @@
 from urllib3.util.ssl_ import create_urllib3_context
 from pathlib import Path
 
-from shared.dataset_service.dataset_service_commons import DatasetTrace
-from shared.helpers.logger import get_logger
-
 
 def create_bucket(bucket_name):
     """
@@ -283,6 +280,8 @@ def record_execution_trace(
     Record the trace in the datastore
     """
     from shared.dataset_service.main import DatasetTraceService
+    from shared.dataset_service.dataset_service_commons import DatasetTrace
+    from shared.helpers.logger import get_logger
 
     trace_service = DatasetTraceService()
 

functions-python/helpers/utils.py

@@ -221,13 +221,18 @@ def extract_location_aggregate_geopolygons(
 ) -> Optional[GeopolygonAggregate]:
     admin_levels = {g.admin_level for g in geopolygons}
     # If duplicates per admin_level exist, resolve instead of returning None
-    if admin_levels != len(geopolygons):
+    if len(admin_levels) != len(geopolygons):
         logger.warning(
             "Duplicate admin levels for point: %s -> %s",
             stop_point,
             geopolygons_as_string(geopolygons),
         )
         geopolygons = dedupe_by_admin_level(geopolygons, logger)
+        logger.warning(
+            "Deduplicated admin levels for point: %s -> %s",
+            stop_point,
+            geopolygons_as_string(geopolygons),
+        )
 
     valid_iso_3166_1 = any(g.iso_3166_1_code for g in geopolygons)
     valid_iso_3166_2 = any(g.iso_3166_2_code for g in geopolygons)

functions-python/reverse_geolocation/src/location_group_utils.py

@@ -80,7 +80,7 @@
     },
 ]
 run_with_feed_index = (
-    5  # Set to an integer index to run with a specific feed from the list above
+    6  # Set to an integer index to run with a specific feed from the list above
 )
 
 

functions-python/reverse_geolocation/src/scripts/reverse_geolocation_process_verifier.py

@@ -106,7 +106,6 @@ def extract_location_aggregates_per_polygon(
             logger.warning("No geopolygons found for point: %s", stop_point)
             continue
 
-        rep_geom = highest.geometry
         country_code = get_country_code_from_polygons(geopolygons)
         if highest.admin_level >= get_country_locality_admin_level(country_code):
             # If admin_level >= locality_admin_level, we can filter points inside this polygon
@@ -144,7 +143,7 @@ def extract_location_aggregates_per_polygon(
 
         # Process ONLY ONE representative point for this stop "cluster"
         location_aggregate = extract_location_aggregate_geopolygons(
-            stop_point=rep_geom,
+            stop_point=stop_point,
             geopolygons=geopolygons,
             logger=logger,
             db_session=db_session,