feat: update update_feed_status to return status diff counts (#992)

Author: sylvansson

functions-python/backfill_dataset_service_date_range/tests/test_backfill_dataset_service_date_range_main.py

@@ -420,7 +420,7 @@ def test_backfill_dataset_service_date_range(mock_backfill_datasets, mock_logger
     with patch.dict(os.environ, {"FEEDS_DATABASE_URL": default_db_url}):
         response_body, status_code = backfill_dataset_service_date_range(None)
 
-    mock_backfill_datasets.asser_called_once()
+    mock_backfill_datasets.assert_called_once()
     assert response_body == "Script executed successfully. 5 datasets updated"
     assert status_code == 200
 
@@ -435,7 +435,7 @@ def test_backfill_dataset_service_date_range_error_raised(
     with patch.dict(os.environ, {"FEEDS_DATABASE_URL": default_db_url}):
         response_body, status_code = backfill_dataset_service_date_range(None)
 
-    mock_backfill_datasets.asser_called_once()
+    mock_backfill_datasets.assert_called_once()
     assert (
         response_body
         == "Error setting the datasets service date range values: Mocked exception"

functions-python/test_utils/database_utils.py

@@ -50,10 +50,10 @@ def get_testing_engine() -> Engine:
     return db._get_engine(echo=False)
 
 
-def get_testing_session() -> Session:
+def get_testing_session(echo: bool = False) -> Session:
     """Returns a SQLAlchemy session for the test db."""
     db = Database(database_url=default_db_url)
-    return db._get_session(echo=False)()
+    return db._get_session(echo=echo)()
 
 
 def clean_testing_db():

functions-python/update_feed_status/src/main.py

@@ -5,7 +5,7 @@
 from shared.helpers.logger import Logger
 from shared.helpers.database import Database
 from typing import TYPE_CHECKING
-from sqlalchemy import case, text
+from sqlalchemy import text
 from shared.database_gen.sqlacodegen_models import Gtfsdataset, Feed, t_feedsearch
 from shared.helpers.database import refresh_materialized_view
 
@@ -33,32 +33,40 @@ def update_feed_statuses_query(session: "Session"):
         .subquery()
     )
 
-    new_status = case(
+    status_conditions = [
         (
             latest_dataset_subq.c.service_date_range_end < today_utc,
-            text("'inactive'::status"),
+            "inactive",
         ),
         (
             latest_dataset_subq.c.service_date_range_start > today_utc,
-            text("'future'::status"),
+            "future",
         ),
         (
             (latest_dataset_subq.c.service_date_range_start <= today_utc)
             & (latest_dataset_subq.c.service_date_range_end >= today_utc),
-            text("'active'::status"),
+            "active",
         ),
-    )
+    ]
 
     try:
-        updated_count = (
-            session.query(Feed)
-            .filter(
-                Feed.status != text("'deprecated'::status"),
-                Feed.status != text("'development'::status"),
-                Feed.id == latest_dataset_subq.c.feed_id,
+        diff_counts: dict[str, int] = {}
+
+        for service_date_conditions, status in status_conditions:
+            diff_counts[status] = (
+                session.query(Feed)
+                .filter(
+                    Feed.id == latest_dataset_subq.c.feed_id,
+                    Feed.status != text("'deprecated'::status"),
+                    Feed.status != text("'development'::status"),
+                    # We filter out feeds that already have the status so that the
+                    # update count reflects the number of feeds that actually
+                    # changed status.
+                    Feed.status != text("'%s'::status" % status),
+                    service_date_conditions,
+                )
+                .update({Feed.status: status}, synchronize_session=False)
             )
-            .update({Feed.status: new_status}, synchronize_session=False)
-        )
     except Exception as e:
         logging.error(f"Error updating feed statuses: {e}")
         raise Exception(f"Error updating feed statuses: {e}")
@@ -68,7 +76,7 @@ def update_feed_statuses_query(session: "Session"):
         refresh_materialized_view(session, t_feedsearch.name)
         logging.info("Feed Database changes committed.")
         session.close()
-        return updated_count
+        return diff_counts
     except Exception as e:
         logging.error("Error committing changes:", e)
         session.rollback()
@@ -81,14 +89,12 @@ def update_feed_status(_):
     """Updates the Feed status based on the latets dataset service date range."""
     Logger.init_logger()
     db = Database(database_url=os.getenv("FEEDS_DATABASE_URL"))
-    update_count = 0
     try:
         with db.start_db_session() as session:
             logging.info("Database session started.")
-            update_count = update_feed_statuses_query(session)
+            diff_counts = update_feed_statuses_query(session)
+            return diff_counts, 200
 
     except Exception as error:
         logging.error(f"Error updating the feed statuses: {error}")
         return f"Error updating the feed statuses: {error}", 500
-
-    return f"Script executed successfully. {update_count} feeds updated", 200

functions-python/update_feed_status/tests/conftest.py

@@ -0,0 +1,89 @@
+#
+#   MobilityData 2025
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#   You may obtain a copy of the License at
+#
+#        http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+#
+
+from datetime import datetime, timedelta
+from uuid import uuid4
+from shared.database_gen.sqlacodegen_models import (
+    Feed,
+    Gtfsdataset,
+)
+from test_shared.test_utils.database_utils import (
+    clean_testing_db,
+    get_testing_session,
+)
+
+future_date = datetime.now() + timedelta(days=15)
+past_date = datetime.now() - timedelta(days=15)
+
+
+def make_dataset(
+    feed_id: str, latest: bool, start: datetime, end: datetime
+) -> Gtfsdataset:
+    return Gtfsdataset(
+        id=str(uuid4()),
+        feed_id=feed_id,
+        latest=latest,
+        service_date_range_start=start,
+        service_date_range_end=end,
+    )
+
+
+def populate_database():
+    session = get_testing_session()
+
+    id_range_by_status = {
+        "inactive": (0, 6),
+        "active": (7, 10),
+        "deprecated": (11, 15),
+        "development": (16, 17),
+        "future": (18, 29),
+    }
+    for status, (a, b) in id_range_by_status.items():
+        for _id in map(str, range(a, b + 1)):
+            session.add(Feed(id=str(_id), status=status))
+
+    # -> inactive
+    for _id in [
+        "0",  # already inactive
+        "7",
+        "8",
+        "22",
+    ]:
+        session.add(make_dataset(_id, True, past_date, past_date))
+
+    # -> active
+    for _id in [
+        "2",
+        "9",  # already active
+        "12",  # deprecated
+        "16",  # development
+        "25",
+    ]:
+        session.add(make_dataset(_id, True, past_date, future_date))
+
+    # -> future
+    for _id in [
+        "10",
+    ]:
+        session.add(make_dataset(_id, False, past_date, future_date))
+        session.add(make_dataset(_id, True, future_date, future_date))
+
+    session.commit()
+
+
+def pytest_sessionstart(session):
+    clean_testing_db()
+    populate_database()

functions-python/update_feed_status/tests/test_update_feed_status_main.py

@@ -1,31 +1,63 @@
 from unittest.mock import patch, MagicMock
-from test_shared.test_utils.database_utils import default_db_url
-from main import update_feed_status, update_feed_statuses_query
+from test_shared.test_utils.database_utils import default_db_url, get_testing_session
+from main import (
+    update_feed_status,
+    update_feed_statuses_query,
+)
+from shared.database_gen.sqlacodegen_models import Feed
 from datetime import date, timedelta
+from typing import Iterator, NamedTuple
 
 import os
 
-
-def test_update_feed_status_return():
-    mock_session = MagicMock()
-
-    today = date(2025, 3, 1)
-
-    mock_subquery = MagicMock()
-    mock_subquery.c.feed_id = 1
-    mock_subquery.c.service_date_range_start = today - timedelta(days=10)
-    mock_subquery.c.service_date_range_end = today + timedelta(days=10)
-
-    mock_query = mock_session.query.return_value
-    mock_query.filter.return_value.subquery.return_value = mock_subquery
-
-    mock_update_query = mock_session.query.return_value.filter.return_value
-    mock_update_query.update.return_value = 3
-
-    updated_count = update_feed_statuses_query(mock_session)
-
-    assert updated_count == 3
-    mock_session.commit.assert_called_once()
+from sqlalchemy import text
+from sqlalchemy.orm import Session
+
+
+class PartialFeed(NamedTuple):
+    """
+    Subset of the Feed entity with only the fields queried in `fetch_feeds`.
+    """
+
+    id: str
+    status: str
+
+
+def fetch_feeds(session: Session) -> Iterator[PartialFeed]:
+    # When adding or removing fields here, `PartialFeed` should be updated to
+    # match, for type safety.
+    query = session.query(Feed.id, Feed.status).filter(
+        Feed.status != text("'deprecated'::status"),
+        Feed.status != text("'development'::status"),
+    )
+    for feed in query:
+        yield PartialFeed(id=feed.id, status=feed.status)
+
+
+def test_update_feed_status():
+    session = get_testing_session()
+    feeds_before: dict[str, PartialFeed] = {f.id: f for f in fetch_feeds(session)}
+    result = dict(update_feed_statuses_query(session))
+    assert result == {
+        "inactive": 3,
+        "active": 2,
+        "future": 1,
+    }
+
+    feeds_after: dict[str, PartialFeed] = {f.id: f for f in fetch_feeds(session)}
+    expected_status_changes = {
+        "2": "active",
+        "7": "inactive",
+        "8": "inactive",
+        "10": "future",
+        "22": "inactive",
+        "25": "active",
+    }
+    for feed_id, feed_before in feeds_before.items():
+        feed_after = feeds_after[feed_id]
+        assert feed_after.status == expected_status_changes.get(
+            feed_id, feed_before.status
+        )
 
 
 def test_update_feed_status_failed_query():
@@ -53,13 +85,14 @@ def test_update_feed_status_failed_query():
 @patch("main.Logger", autospec=True)
 @patch("main.update_feed_statuses_query")
 def test_updated_feed_status(mock_update_query, mock_logger):
-    mock_update_query.return_value = 5
+    return_value = {"active": 5}
+    mock_update_query.return_value = return_value
 
     with patch.dict(os.environ, {"FEEDS_DATABASE_URL": default_db_url}):
         response_body, status_code = update_feed_status(None)
 
-    mock_update_query.asser_called_once()
-    assert response_body == "Script executed successfully. 5 feeds updated"
+    mock_update_query.assert_called_once()
+    assert response_body == return_value
     assert status_code == 200
 
 
@@ -71,6 +104,6 @@ def test_updated_feed_status_error_raised(mock_update_query, mock_logger):
     with patch.dict(os.environ, {"FEEDS_DATABASE_URL": default_db_url}):
         response_body, status_code = update_feed_status(None)
 
-    mock_update_query.asser_called_once()
+    mock_update_query.assert_called_once()
     assert response_body == "Error updating the feed statuses: Mocked exception"
     assert status_code == 500