reuse common trigger_download_dataset function

Author: davidgamez

functions-python/helpers/pub_sub.py

@@ -23,7 +23,7 @@
 from google.cloud.pubsub_v1 import PublisherClient
 from google.cloud.pubsub_v1.publisher.futures import Future
 
-from shared.database_gen.sqlacodegen_models import Feed
+from shared.database_gen.sqlacodegen_models import Feed, Gtfsfeed
 
 PROJECT_ID = os.getenv("PROJECT_ID")
 DATASET_BATCH_TOPIC = os.getenv("DATASET_PROCESSING_TOPIC_NAME")
@@ -73,7 +73,7 @@ def publish_messages(data: List[Dict], project_id, topic_name) -> None:
 
 
 def trigger_dataset_download(
-    feed: Feed,
+    feed: Feed | Gtfsfeed,
     execution_id: str,
     topic_name: str = DATASET_BATCH_TOPIC,
 ) -> None:

functions-python/operations_api/src/feeds_operations/impl/feeds_operations_impl.py

@@ -13,7 +13,6 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 #
-import os
 
 from fastapi.encoders import jsonable_encoder
 from fastapi.responses import JSONResponse
@@ -57,7 +56,7 @@
     Feed,
     Gtfsrealtimefeed,
 )
-from shared.helpers.pub_sub import get_execution_id, publish_messages
+from shared.helpers.pub_sub import get_execution_id, trigger_dataset_download
 from shared.helpers.query_helper import (
     query_feed_by_stable_id,
     get_feeds_query,
@@ -74,9 +73,6 @@
 from .models.operation_gtfs_rt_feed_impl import OperationGtfsRtFeedImpl
 from .request_validator import validate_request
 
-pubsub_topic_name = os.getenv("DATASET_PROCESSING_TOPIC_NAME")
-project_id = os.getenv("PROJECT_ID")
-
 
 class OperationsApiImpl(BaseOperationsApi):
     """Implementation of the operations API."""
@@ -112,24 +108,24 @@ def assert_no_existing_feed_url(producer_url: str, db_session: Session):
                 detail=message,
             )
 
-    @staticmethod
-    def send_feed_process_event(feed: type[Gtfsfeed] | None, request=None):
-        """Send a message to Pub/Sub to process the feed."""
-        message_payload = {
-            "execution_id": get_execution_id(
-                get_request_context(), "feed-created-process"
-            ),
-            "producer_url": feed.producer_url,
-            "feed_stable_id": feed.stable_id,
-            "feed_id": feed.id,
-            "dataset_stable_id": None,
-            "dataset_hash": None,
-            "authentication_type": feed.authentication_type,
-            "authentication_info_url": feed.authentication_info_url,
-            "api_key_parameter_name": feed.api_key_parameter_name,
-        }
-        publish_messages([message_payload], project_id, pubsub_topic_name)
-        logging.debug("Sent feed process event")
+    # @staticmethod
+    # def send_feed_process_event(feed: type[Gtfsfeed] | None, request=None):
+    #     """Send a message to Pub/Sub to process the feed."""
+    #     message_payload = {
+    #         "execution_id": get_execution_id(
+    #             get_request_context(), "feed-created-process"
+    #         ),
+    #         "producer_url": feed.producer_url,
+    #         "feed_stable_id": feed.stable_id,
+    #         "feed_id": feed.id,
+    #         "dataset_stable_id": None,
+    #         "dataset_hash": None,
+    #         "authentication_type": feed.authentication_type,
+    #         "authentication_info_url": feed.authentication_info_url,
+    #         "api_key_parameter_name": feed.api_key_parameter_name,
+    #     }
+    #     publish_messages([message_payload], project_id, pubsub_topic_name)
+    #     logging.debug("Sent feed process event")
 
     @with_db_session
     async def get_feeds(
@@ -391,7 +387,10 @@ async def create_gtfs_feed(
         db_session.add(new_feed)
         db_session.commit()
         created_feed = db_session.get(Gtfsfeed, new_feed.id)
-        self.send_feed_process_event(created_feed)
+        trigger_dataset_download(
+            created_feed,
+            get_execution_id(get_request_context(), "feed-created-process"),
+        )
         logging.info("Created new GTFS feed with ID: %s", new_feed.stable_id)
         payload = OperationGtfsFeedImpl.from_orm(created_feed).model_dump()
         return JSONResponse(status_code=201, content=jsonable_encoder(payload))

functions-python/operations_api/tests/feeds_operations/impl/test_create_feeds_operations_impl_gtfs.py

@@ -24,9 +24,6 @@
 import uuid
 from unittest.mock import patch
 
-# For asserting module-level topic/project used on publish
-from feeds_operations.impl import feeds_operations_impl as ops_module
-
 
 @pytest.fixture
 def update_request_gtfs_feed():
@@ -65,7 +62,7 @@ def db_session():
 
 
 @pytest.mark.asyncio
-@patch("feeds_operations.impl.feeds_operations_impl.publish_messages")
+@patch("feeds_operations.impl.feeds_operations_impl.trigger_dataset_download")
 async def test_create_gtfs_feed_success(mock_publish_messages, db_session):
     api = OperationsApiImpl()
     unique_url = f"https://new-feed.example.com/{uuid.uuid4()}"
@@ -114,27 +111,20 @@ async def test_create_gtfs_feed_success(mock_publish_messages, db_session):
         # Assert publish_messages was called exactly once with expected payload
         assert mock_publish_messages.call_count == 1
         args, kwargs = mock_publish_messages.call_args
-        assert len(args) == 3  # data list, project_id, topic_name
-        data_list, project_id, topic_name = args
-
-        # Project/topic should be whatever module-level variables resolve to (may be None in tests)
-        assert project_id == ops_module.project_id
-        assert topic_name == ops_module.pubsub_topic_name
+        assert len(args) == 2  # data list, project_id, topic_name
+        feed, execution_id = args
 
         # Validate message payload shape and values
-        assert isinstance(data_list, list) and len(data_list) == 1
-        message = data_list[0]
-        assert message["producer_url"] == unique_url
-        assert message["feed_stable_id"] == payload["stable_id"]
-        assert message["feed_id"] == payload["id"]
-        assert message["dataset_stable_id"] is None
-        assert message["dataset_hash"] is None
-        assert message["authentication_type"] == "0"
-        assert message["authentication_info_url"] is None
-        assert message["api_key_parameter_name"] is None
+        # assert isinstance(feed, list) and len(data_list) == 1
+        # message = feed[0]
+        assert feed.producer_url == unique_url
+        assert feed.stable_id == payload["stable_id"]
+        assert feed.id == payload["id"]
+        assert feed.authentication_type == "0"
+        assert feed.authentication_info_url is None
+        assert feed.api_key_parameter_name is None
         # Non-deterministic but must start with expected prefix
-        assert isinstance(message.get("execution_id"), str)
-        assert message["execution_id"].startswith("feed-created-process-")
+        assert execution_id.startswith("feed-created-process-")
     finally:
         # Cleanup to avoid impacting other tests
         stable_id = payload.get("stable_id") if isinstance(payload, dict) else None