fixed broken tests

functions-python/feed_sync_process_transitland/src/main.py

@@ -27,7 +27,7 @@
 from database_gen.sqlacodegen_models import Feed, Externalid, Redirectingid
 from sqlalchemy.exc import SQLAlchemyError
 
-from helpers.database import start_db_session, close_db_session
+from helpers.database import Database, with_db_session
 from helpers.logger import Logger, StableIdFilter
 from helpers.feed_sync.models import TransitFeedSyncPayload as FeedPayload
 from helpers.locations import create_or_get_location
@@ -455,21 +455,15 @@ def process_feed_event(cloud_event):
         # Decode payload from Pub/Sub message
         pubsub_message = base64.b64decode(cloud_event.data["message"]["data"]).decode()
         message_data = json.loads(pubsub_message)
-
         payload = FeedPayload(**message_data)
-
-        db_session = start_db_session(FEEDS_DATABASE_URL)
-
-        try:
+        db = Database(FEEDS_DATABASE_URL)
+        with db.start_db_session() as db_session:
             processor = FeedProcessor(db_session)
             processor.process_feed(payload)
 
             log_message("info", f"Successfully processed feed: {payload.external_id}")
             return "Success", 200
 
-        finally:
-            close_db_session(db_session)
-
     except Exception as e:
         error_msg = f"Error processing feed event: {str(e)}"
         log_message("error", error_msg)

functions-python/feed_sync_process_transitland/tests/test_feed_sync_process.py

@@ -43,6 +43,12 @@ def mock_location():
     return Mock()
 
 
+@pytest.fixture
+def mock_db():
+    with patch("feed_sync_process_transitland.src.main.Database") as mock_db:
+        yield mock_db
+
+
 class MockLogger:
     """Mock logger for testing"""
 
@@ -608,7 +614,7 @@ def __init__(self):
             assert not mock_new_feed.locations
 
     def test_process_feed_event_database_connection_error(
-        self, processor, feed_payload, mock_logging
+        self, processor, feed_payload, mock_logging, mock_db
     ):
         """Test feed event processing with database connection error."""
         # Create cloud event with valid payload
@@ -620,21 +626,18 @@ def test_process_feed_event_database_connection_error(
         cloud_event.data = {"message": {"data": payload_data}}
 
         # Mock database session to raise error
-        with patch(
-            "feed_sync_process_transitland.src.main.start_db_session"
-        ) as mock_start_session:
-            mock_start_session.side_effect = SQLAlchemyError(
-                "Database connection error"
-            )
+        mock_db.return_value.start_db_session.side_effect = SQLAlchemyError(
+            "Database connection error"
+        )
 
-            result = process_feed_event(cloud_event)
-            assert result[1] == 500
-            mock_logging.error.assert_called_with(
-                "Error processing feed event: Database connection error"
-            )
+        result = process_feed_event(cloud_event)
+        assert result[1] == 500
+        mock_logging.error.assert_called_with(
+            "Error processing feed event: Database connection error"
+        )
 
     def test_process_feed_event_pubsub_error(
-        self, processor, feed_payload, mock_logging
+        self, processor, feed_payload, mock_logging, mock_db
     ):
         """Test feed event processing handles missing credentials error."""
         # Create cloud event with valid payload
@@ -648,19 +651,18 @@ def test_process_feed_event_pubsub_error(
         cloud_event.data = {"message": {"data": payload_data}}
 
         # Mock database session with minimal setup
-        mock_session = Mock()
+        mock_session = MagicMock()
         mock_session.query.return_value.filter.return_value.all.return_value = []
+        mock_db.return_value.start_db_session.return_value.__enter__.return_value = (
+            mock_session
+        )
 
         # Process event and verify error handling
-        with patch(
-            "feed_sync_process_transitland.src.main.start_db_session",
-            return_value=mock_session,
-        ):
-            result = process_feed_event(cloud_event)
-            assert result[1] == 500
-            mock_logging.error.assert_called_with(
-                "Error processing feed event: File dummy-credentials.json was not found."
-            )
+        result = process_feed_event(cloud_event)
+        assert result[1] == 500
+        mock_logging.error.assert_called_with(
+            "Error processing feed event: File dummy-credentials.json was not found."
+        )
 
     def test_process_feed_event_malformed_cloud_event(self, mock_logging):
         """Test feed event processing with malformed cloud event."""