Add return to main function and db_session commit

Author: davidgamez

functions-python/process_validation_report/tests/conftest.py

@@ -143,7 +143,7 @@ def pytest_sessionfinish(session, exitstatus):
     returning the exit status to the system.
     """
     # Cleaned at the beginning instead of the end so we can examine the DB after the test.
-    # clean_testing_db()
+    clean_testing_db()
 
 
 def pytest_unconfigure(config):

functions-python/process_validation_report/tests/test_validation_report.py

@@ -82,6 +82,7 @@ def test_get_dataset(self, db_session):
         )
         try:
             db_session.add(feed)
+            db_session.flush()
             db_session.add(dataset)
             db_session.flush()
             returned_dataset = get_dataset(dataset_stable_id, db_session)
@@ -123,6 +124,7 @@ def test_create_validation_report_entities(self, mock_get, db_session):
         )
         try:
             db_session.add(feed)
+            db_session.flush()
             db_session.add(dataset)
             db_session.commit()
             create_validation_report_entities(feed_stable_id, dataset_stable_id, "1.0")
@@ -347,17 +349,18 @@ def test_create_validation_report_entities_missing_validator_version(
                 ],
             },
         )
-        feed_stable_id = faker.word()
-        dataset_stable_id = faker.word()
+        feed_stable_id = faker.uuid4()
+        dataset_stable_id = faker.uuid4()
 
         # Create GTFS Feed
-        feed = Gtfsfeed(id=faker.word(), data_type="gtfs", stable_id=feed_stable_id)
+        feed = Gtfsfeed(id=faker.uuid4(), data_type="gtfs", stable_id=feed_stable_id)
         # Create a new dataset
         dataset = Gtfsdataset(
-            id=faker.word(), feed_id=feed.id, stable_id=dataset_stable_id, latest=True
+            id=faker.uuid4(), feed_id=feed.id, stable_id=dataset_stable_id, latest=True
         )
         try:
             db_session.add(feed)
+            db_session.flush()
             db_session.add(dataset)
             db_session.commit()
             create_validation_report_entities(feed_stable_id, dataset_stable_id, "1.0")

functions-python/tasks_executor/src/tasks/geojson/update_geojson_files_precision.py

@@ -71,15 +71,19 @@ def _upload_file(bucket, geojson):
 
 
 @track_metrics(metrics=("time", "memory", "cpu"))
-def _update_feed_info(feed, timestamp):
+def _update_feed_info(feed: Gtfsfeed, timestamp):
     feed.geolocation_file_created_date = timestamp
     # find the most recent dataset with bounding box and set the id
     if feed.gtfsdatasets and any(d.bounding_box for d in feed.gtfsdatasets):
         latest_with_bbox = max(
             (d for d in feed.gtfsdatasets if d.bounding_box),
-            key=lambda d: d.downloaded_date or timestamp,
+            key=lambda d: d.downloaded_at or timestamp,
+        )
+        feed.geolocation_file_dataset_id = latest_with_bbox.id
+    else:
+        logging.info(
+            "No GTFS datasets available with bounding box for feed %s", feed.id
         )
-        feed.geolocation_file_dataset_id = latest_with_bbox.bounding_box.id
 
 
 @track_metrics(metrics=("time", "memory", "cpu"))
@@ -171,7 +175,9 @@ def update_geojson_files_precision_handler(
 
             geojson = process_geojson(geojson, precision)
             if not geojson:
-                logging.info("No valid GeoJSON features found in %s", file.name)
+                logging.warning("No valid GeoJSON features found in %s", file.name)
+                errors.append(feed.stable_id)
+                continue
 
             # Optionally upload processed geojson
             if not dry_run:
@@ -182,7 +188,8 @@ def update_geojson_files_precision_handler(
         except Exception as e:
             logging.exception("Error processing feed %s: %s", feed.stable_id, e)
             errors.append(feed.stable_id)
-
+    if not dry_run and processed > 0:
+        db_session.commit()
     summary = {
         "total_processed_files": processed,
         "errors": errors,
@@ -194,4 +201,4 @@ def update_geojson_files_precision_handler(
         },
     }
     logging.info("update_geojson_files_handler result: %s", summary)
-    return
+    return summary

functions-python/tasks_executor/tests/conftest.py

@@ -35,6 +35,7 @@ def populate_database(db_session: Session | None = None):
     """
     # Create 2 GTFS Feeds
     feeds = []
+    # raise NotImplementedError("Implement the function to populate the database.")
     now = datetime.now(UTC)
     for i in range(2):
         feed = Gtfsfeed(
@@ -56,6 +57,7 @@ def populate_database(db_session: Session | None = None):
             feed=feed,
             stable_id=f"dataset_stable_{i:04d}",
             downloaded_at=now - timedelta(days=i),
+            bounding_box="POLYGON((0 0, 0 1, 1 1, 1 0, 0 0))",
         )
         db_session.add(dataset)
         datasets.append(dataset)

functions-python/tasks_executor/tests/tasks/geojson/test_update_geojson_files_precision.py

@@ -5,26 +5,32 @@
 import unittest
 from unittest.mock import patch
 
+from sqlalchemy.orm import Session
+
+from shared.database.database import with_db_session
+from shared.helpers.src.shared.database_gen.sqlacodegen_models import Gtfsfeed
 from tasks.geojson.update_geojson_files_precision import (
     process_geojson,
     update_geojson_files_precision_handler,
     GEOLOCATION_FILENAME,
 )
+from test_shared.test_utils.database_utils import default_db_url
 
 
 class _FakeBlobContext:
-    def __init__(self, bucket, name):
+    def __init__(self, bucket, name, blob_exists=True):
         self.bucket = bucket
         self.name = name
+        self.blob_exists = blob_exists
 
     def __enter__(self):
         return self
 
-    def __exit__(self, exc_type, exc, tb):
-        return False
+    # def __exit__(self, exc_type, exc, tb):
+    #     return self.exists
 
     def exists(self):
-        return self.name in self.bucket.initial_blobs
+        return self.blob_exists
 
     def download_as_text(self):
         return self.bucket.initial_blobs[self.name]
@@ -60,15 +66,16 @@ def bucket(self, name):
 
 
 class FakeStorageModule:
-    def __init__(self, bucket):
+    def __init__(self, bucket, blob_exists=True):
         self._bucket = bucket
+        self._blob_exists = blob_exists
 
     def Client(self):
         return FakeClient(self._bucket)
 
     # storage.Blob(...) used as a context manager in the handler
     def Blob(self, *, bucket, name):
-        return _FakeBlobContext(bucket, name)
+        return _FakeBlobContext(bucket, name, self._blob_exists)
 
 
 class TestUpdateGeojsonFilesPrecision(unittest.TestCase):
@@ -129,8 +136,8 @@ def test_process_geojson_single_feature_and_list_variants(self):
             [round(1.23456789, 3), round(2.3456789, 3)],
         )
 
-    @patch("tasks.geojson.update_geojson_files_precision.query_unprocessed_feeds")
-    def test_handler_uploads_and_updates_feed_info(self, mock_query):
+    @with_db_session(db_url=default_db_url)
+    def test_handler_uploads_and_updates_feed_info(self, db_session: Session):
         geo = {
             "type": "FeatureCollection",
             "features": [
@@ -144,11 +151,12 @@ def test_handler_uploads_and_updates_feed_info(self, mock_query):
                 }
             ],
         }
-        feed_stable_id = "feed_123"
+        testing_feed = db_session.query(Gtfsfeed).limit(1).first()
+        feed_stable_id = testing_feed.stable_id
         blob_name = f"{feed_stable_id}/{GEOLOCATION_FILENAME}"
 
         fake_bucket = FakeBucket(initial_blobs={blob_name: json.dumps(geo)})
-        fake_storage = FakeStorageModule(fake_bucket)
+        fake_storage = FakeStorageModule(fake_bucket, blob_exists=True)
 
         # create module objects for google and google.cloud and inject via sys.modules
         cloud_mod = types.ModuleType("google.cloud")
@@ -157,32 +165,6 @@ def test_handler_uploads_and_updates_feed_info(self, mock_query):
         google_mod = types.ModuleType("google")
         google_mod.cloud = cloud_mod
 
-        fake_feed = types.SimpleNamespace(
-            stable_id=feed_stable_id,
-            geolocation_file_created_date=None,
-            gtfsdatasets=[
-                types.SimpleNamespace(
-                    bounding_box=types.SimpleNamespace(id="bbid"), downloaded_date=None
-                )
-            ],
-        )
-
-        mock_query.return_value = [fake_feed]
-
-        # fake db session
-        class FakeExecResult:
-            def scalar(self):
-                return "NOW_TS"
-
-        class FakeDBSession:
-            def execute(self, q):
-                return FakeExecResult()
-
-            def commit(self):
-                return None
-
-        fake_db = FakeDBSession()
-
         payload = {
             "bucket_name": "any-bucket",
             "dry_run": False,
@@ -193,7 +175,9 @@ def commit(self):
         # Inject modules into sys.modules for the duration of the handler call
         with patch.dict(sys.modules, {"google.cloud": cloud_mod, "google": google_mod}):
             # call wrapped handler to provide fake db_session
-            update_geojson_files_precision_handler(payload, db_session=fake_db)
+            result = update_geojson_files_precision_handler(
+                payload, db_session=db_session
+            )
 
         # verify upload happened
         self.assertIn("geolocation.geojson", fake_bucket.uploaded)
@@ -202,8 +186,111 @@ def commit(self):
         coords = uploaded_geo.get("features")[0]["geometry"]["coordinates"]
         self.assertEqual(coords, [round(100.1234567, 5), round(0.9876543, 5)])
 
+        self.assertEqual(
+            {
+                "total_processed_files": 1,
+                "errors": [],
+                "not_found_file": 0,
+                "params": {
+                    "dry_run": False,
+                    "precision": 5,
+                    "limit": 1,
+                },
+            },
+            result,
+        )
         # feed updated
-        self.assertEqual(fake_feed.geolocation_file_created_date, "NOW_TS")
+        reloaded_testing_feed = (
+            db_session.query(Gtfsfeed)
+            .filter(Gtfsfeed.id.__eq__(testing_feed.id))
+            .limit(1)
+            .first()
+        )
+        self.assertIsNotNone(reloaded_testing_feed.geolocation_file_dataset_id)
+        self.assertIsNotNone(reloaded_testing_feed.geolocation_file_created_date)
+
+    @with_db_session(db_url=default_db_url)
+    def test_handler_file_dont_exists(self, db_session: Session):
+        fake_bucket = FakeBucket(initial_blobs={})
+        fake_storage = FakeStorageModule(fake_bucket, blob_exists=False)
+
+        # create module objects for google and google.cloud and inject via sys.modules
+        cloud_mod = types.ModuleType("google.cloud")
+        # 'from google.cloud import storage' in handler will bind 'storage' to this attribute
+        cloud_mod.storage = fake_storage
+        google_mod = types.ModuleType("google")
+        google_mod.cloud = cloud_mod
+
+        payload = {
+            "bucket_name": "any-bucket",
+            "dry_run": False,
+            "precision": 5,
+            "limit": 1,
+        }
+
+        # Inject modules into sys.modules for the duration of the handler call
+        with patch.dict(sys.modules, {"google.cloud": cloud_mod, "google": google_mod}):
+            # call wrapped handler to provide fake db_session
+            result = update_geojson_files_precision_handler(
+                payload, db_session=db_session
+            )
+            self.assertEqual(
+                {
+                    "total_processed_files": 0,
+                    "errors": [],
+                    "not_found_file": 1,
+                    "params": {
+                        "dry_run": False,
+                        "precision": 5,
+                        "limit": 1,
+                    },
+                },
+                result,
+            )
+
+    @with_db_session(db_url=default_db_url)
+    def test_handler_file_not_valid_file(self, db_session: Session):
+        geo = "{}"
+        testing_feed = db_session.query(Gtfsfeed).limit(1).first()
+        feed_stable_id = testing_feed.stable_id
+        blob_name = f"{feed_stable_id}/{GEOLOCATION_FILENAME}"
+
+        fake_bucket = FakeBucket(initial_blobs={blob_name: geo})
+        fake_storage = FakeStorageModule(fake_bucket, blob_exists=True)
+
+        # create module objects for google and google.cloud and inject via sys.modules
+        cloud_mod = types.ModuleType("google.cloud")
+        # 'from google.cloud import storage' in handler will bind 'storage' to this attribute
+        cloud_mod.storage = fake_storage
+        google_mod = types.ModuleType("google")
+        google_mod.cloud = cloud_mod
+
+        payload = {
+            "bucket_name": "any-bucket",
+            "dry_run": False,
+            "precision": 5,
+            "limit": 1,
+        }
+        testing_feed = db_session.query(Gtfsfeed).limit(1).first()
+        # Inject modules into sys.modules for the duration of the handler call
+        with patch.dict(sys.modules, {"google.cloud": cloud_mod, "google": google_mod}):
+            # call wrapped handler to provide fake db_session
+            result = update_geojson_files_precision_handler(
+                payload, db_session=db_session
+            )
+            self.assertEqual(
+                {
+                    "total_processed_files": 0,
+                    "errors": [testing_feed.stable_id],
+                    "not_found_file": 0,
+                    "params": {
+                        "dry_run": False,
+                        "precision": 5,
+                        "limit": 1,
+                    },
+                },
+                result,
+            )
 
 
 if __name__ == "__main__":