diff --git a/mbta-performance/.chalice/config.json b/mbta-performance/.chalice/config.json index f28777f..0a78489 100644 --- a/mbta-performance/.chalice/config.json +++ b/mbta-performance/.chalice/config.json @@ -25,7 +25,17 @@ }, "process_yesterday_lamp": { "iam_policy_file": "policy-lamp-ingest.json", - "lambda_timeout": 900, + "lambda_timeout": 600, + "lambda_memory_size": 4096 + }, + "process_daily_bus_lamp": { + "iam_policy_file": "policy-bus-lamp-ingest.json", + "lambda_timeout": 120, + "lambda_memory_size": 2048 + }, + "process_yesterday_bus_lamp": { + "iam_policy_file": "policy-bus-lamp-ingest.json", + "lambda_timeout": 600, "lambda_memory_size": 4096 }, "regenerate_tm_benchmarks": { @@ -36,4 +46,4 @@ } } } -} +} \ No newline at end of file diff --git a/mbta-performance/.chalice/policy-bus-lamp-ingest.json b/mbta-performance/.chalice/policy-bus-lamp-ingest.json new file mode 100644 index 0000000..2259436 --- /dev/null +++ b/mbta-performance/.chalice/policy-bus-lamp-ingest.json @@ -0,0 +1,19 @@ +{ + "Version": "2012-10-17", + "Statement": [ + { + "Action": [ + "logs:CreateLogGroup", + "logs:CreateLogStream", + "logs:PutLogEvents" + ], + "Effect": "Allow", + "Resource": "arn:*:logs:*:*:*" + }, + { + "Action": ["s3:GetObject", "s3:PutObject"], + "Effect": "Allow", + "Resource": ["arn:aws:s3:::tm-mbta-performance/Events-lamp/bus-daily-data/*"] + } + ] +} diff --git a/mbta-performance/app.py b/mbta-performance/app.py index 50aeedd..5b601fd 100644 --- a/mbta-performance/app.py +++ b/mbta-performance/app.py @@ -46,3 +46,23 @@ def process_yesterday_lamp(event): def regenerate_tm_benchmarks(event): """Regenerate TransitMatters travel-time benchmarks for rapid transit.""" benchmarks.generate_travel_time_benchmarks() + + +# Bus LAMP data processing +# Runs every 30 minutes from either 5 AM -> 2:30AM or 6 AM -> 3:30 AM depending on DST +@app.schedule(Cron("*/30", "0-7,10-23", "*", "*", "?", "*")) +def process_daily_bus_lamp(event): + """Ingest today's bus LAMP data.""" + now_boston = datetime.now(ZoneInfo("US/Eastern")) + + if now_boston.hour >= 3 and now_boston.hour < 6: + return + + lamp.ingest_today_bus_data() + + +# Runs once the next day at 11am or 12pm depending on DST +@app.schedule(Cron("0", "15", "*", "*", "?", "*")) +def process_yesterday_bus_lamp(event): + """Process yesterday's bus LAMP data, to ensure we have everything we need.""" + lamp.ingest_yesterday_bus_data() diff --git a/mbta-performance/chalicelib/lamp/__init__.py b/mbta-performance/chalicelib/lamp/__init__.py index d6a8f0d..aeecb47 100644 --- a/mbta-performance/chalicelib/lamp/__init__.py +++ b/mbta-performance/chalicelib/lamp/__init__.py @@ -1,3 +1,9 @@ -__all__ = ["ingest_today_lamp_data", "ingest_yesterday_lamp_data"] +__all__ = [ + "ingest_today_lamp_data", + "ingest_yesterday_lamp_data", + "ingest_today_bus_data", + "ingest_yesterday_bus_data", +] +from .bus_ingest import ingest_today_bus_data, ingest_yesterday_bus_data from .ingest import ingest_today_lamp_data, ingest_yesterday_lamp_data diff --git a/mbta-performance/chalicelib/lamp/backfill/bus.py b/mbta-performance/chalicelib/lamp/backfill/bus.py new file mode 100644 index 0000000..923c363 --- /dev/null +++ b/mbta-performance/chalicelib/lamp/backfill/bus.py @@ -0,0 +1,76 @@ +import logging +import os +from datetime import date, timedelta + +import pandas as pd + +from ... import parallel +from ..bus_ingest import RTE_DIR_STOP, fetch_bus_pq_file_from_remote, ingest_bus_pq_file, upload_bus_to_s3 + +logger = logging.getLogger(__name__) + +_parallel_upload = parallel.make_parallel(upload_bus_to_s3) + +EARLIEST_BUS_LAMP_DATA = date(2020, 1, 8) + +LOCAL_ARCHIVE_PATH = os.environ.get("LOCAL_ARCHIVE_PATH", "./feeds") + + +def backfill_all_bus_dates(start_date: date = EARLIEST_BUS_LAMP_DATA): + """Backfill all dates with bus LAMP data, most recent to oldest.""" + + if start_date < EARLIEST_BUS_LAMP_DATA: + raise ValueError(f"start_date {start_date} is before earliest available data {EARLIEST_BUS_LAMP_DATA}") + + end_date = date.today() - timedelta(days=1) + dates = pd.date_range(start_date, end_date, freq="d") + num_days = len(dates) + estimated_cost = num_days / 30 + + confirmation = input( + f"This will backfill {num_days} days of bus LAMP data ({start_date} to {end_date}). " + f"Estimated cost: ~${estimated_cost:.2f} (based on ~$1 per 30 days). " + "This will take hours. Are you sure you want to proceed? (yes/no): " + ) + if confirmation.lower() not in ("yes", "y"): + print("You must enter 'yes' to proceed. Exiting.") + exit(1) + + for backfill_timestamp in dates[::-1]: + date_to_backfill = backfill_timestamp.date() + try: + pq_df = fetch_bus_pq_file_from_remote(date_to_backfill) + except ValueError as e: + logger.warning(f"Failed to fetch {date_to_backfill}: {e}") + continue + + logger.info(f"Processing {date_to_backfill}") + processed = ingest_bus_pq_file(pq_df, date_to_backfill, local_archive_path=LOCAL_ARCHIVE_PATH) + + group_event_groups = processed.groupby(RTE_DIR_STOP) + logger.info(f"Uploading events for {len(group_event_groups)} route-direction-stop groups to S3") + _parallel_upload(group_event_groups, date_to_backfill) + logger.info(f"Finished {date_to_backfill}") + + +if __name__ == "__main__": + import argparse + import os + from datetime import datetime + + parser = argparse.ArgumentParser(description="Backfill bus LAMP data from a start date up to yesterday.") + parser.add_argument( + "--start-date", + type=lambda s: datetime.strptime(s, "%Y-%m-%d").date(), + default=EARLIEST_BUS_LAMP_DATA, + help=f"Start date in YYYY-MM-DD format (default: {EARLIEST_BUS_LAMP_DATA}).", + ) + args = parser.parse_args() + + log_level = os.environ.get("LOG_LEVEL", "INFO").upper() + logging.basicConfig( + level=getattr(logging, log_level, logging.INFO), + format="%(asctime)s - %(levelname)s: %(message)s", + datefmt="%H:%M:%S", + ) + backfill_all_bus_dates(start_date=args.start_date) diff --git a/mbta-performance/chalicelib/lamp/bus_constants.py b/mbta-performance/chalicelib/lamp/bus_constants.py new file mode 100644 index 0000000..48cc6ca --- /dev/null +++ b/mbta-performance/chalicelib/lamp/bus_constants.py @@ -0,0 +1,56 @@ +# Remote URL for fetching bus LAMP data (daily parquet files) +BUS_DAILY_URL_TEMPLATE = "https://performancedata.mbta.com/lamp/bus_vehicle_events/{YYYYMMDD}.parquet" + +# Columns to read from the source parquet files +BUS_LAMP_COLUMNS = [ + "service_date", + "route_id", + "trip_id", + "stop_id", + "direction_id", + "stop_sequence", + "vehicle_label", + "previous_stop_id", + # Actual timestamps (UTC-aware) + "stop_arrival_dt", + "stop_departure_dt", + # Actual metrics + "travel_time_seconds", + "stopped_duration_seconds", + "route_direction_headway_seconds", + # Scheduled metrics + "plan_travel_time_seconds", + "plan_route_direction_headway_seconds", +] + +# Columns output to S3 events.csv +BUS_S3_COLUMNS = [ + "service_date", + "route_id", + "trip_id", + "direction_id", + "stop_id", + "stop_sequence", + "vehicle_label", + "event_type", + "event_time", + "travel_time_seconds", + "dwell_time_seconds", + "headway_seconds", + "scheduled_tt", + "scheduled_headway", +] + +BUS_COLUMN_RENAME_MAP = { + "stopped_duration_seconds": "dwell_time_seconds", + "route_direction_headway_seconds": "headway_seconds", + "plan_travel_time_seconds": "scheduled_tt", + "plan_route_direction_headway_seconds": "scheduled_headway", +} + +# Output S3 bucket and key template +BUS_S3_BUCKET = "tm-mbta-performance" +# month and day are not zero-padded +BUS_S3_KEY_TEMPLATE = ( + "Events-lamp/bus-daily-data/{route_id}-{direction_id}-{stop_id}/Year={YYYY}/Month={_M}/Day={_D}/events.csv" +) diff --git a/mbta-performance/chalicelib/lamp/bus_ingest.py b/mbta-performance/chalicelib/lamp/bus_ingest.py new file mode 100644 index 0000000..89ffa16 --- /dev/null +++ b/mbta-performance/chalicelib/lamp/bus_ingest.py @@ -0,0 +1,288 @@ +import io +import logging +from datetime import date +from typing import Tuple + +import pandas as pd +import requests + +from .. import parallel, s3 +from ..date import EASTERN_TIME, get_current_service_date +from ..gtfs import fetch_stop_times_from_gtfs +from .bus_constants import ( + BUS_COLUMN_RENAME_MAP, + BUS_DAILY_URL_TEMPLATE, + BUS_LAMP_COLUMNS, + BUS_S3_BUCKET, + BUS_S3_COLUMNS, + BUS_S3_KEY_TEMPLATE, +) + +logger = logging.getLogger(__name__) + +RTE_DIR_STOP = ["route_id", "direction_id", "stop_id"] + + +def fetch_bus_pq_file_from_remote(service_date: date) -> pd.DataFrame: + """Fetch a bus parquet file from LAMP for a given service date.""" + url = BUS_DAILY_URL_TEMPLATE.format(YYYYMMDD=service_date.strftime("%Y%m%d")) + logger.info(f"Fetching bus LAMP parquet file from {url}") + result = requests.get(url) + + if result.status_code != 200: + logger.error(f"Failed to fetch bus LAMP parquet file from {url}. Status code: {result.status_code}") + raise ValueError(f"Failed to fetch bus LAMP parquet file from {url}. Status code: {result.status_code}") + + logger.info(f"Successfully fetched bus LAMP parquet file ({len(result.content)} bytes)") + df = pd.read_parquet( + io.BytesIO(result.content), + columns=BUS_LAMP_COLUMNS, + engine="pyarrow", + dtype_backend="numpy_nullable", + ) + logger.info(f"Parsed parquet file: {len(df)} rows") + return df + + +def _process_bus_arrival_departure_times(df: pd.DataFrame) -> pd.DataFrame: + """Split bus events into separate ARR and DEP rows. + + Bus LAMP data already provides stop_arrival_dt and stop_departure_dt + as UTC-aware timestamps, along with previous_stop_id. We just need to: + 1. Convert to Eastern Time + 2. Create separate ARR/DEP event rows + 3. For DEP events, use previous_stop_id as the stop_id + """ + logger.debug(f"Processing arrival/departure times for {len(df)} rows") + + # Arrivals: use stop_arrival_dt and the current stop_id + arr_df = df[df["stop_arrival_dt"].notna()].copy() + arr_df["event_type"] = "ARR" + arr_df["event_time"] = arr_df["stop_arrival_dt"].dt.tz_convert(EASTERN_TIME) + arr_df = arr_df[BUS_S3_COLUMNS] + + # Departures: use stop_departure_dt and previous_stop_id + dep_df = df[df["stop_departure_dt"].notna() & df["previous_stop_id"].notna()].copy() + dep_df["event_type"] = "DEP" + dep_df["event_time"] = dep_df["stop_departure_dt"].dt.tz_convert(EASTERN_TIME) + dep_df["stop_id"] = dep_df["previous_stop_id"] + dep_df = dep_df[BUS_S3_COLUMNS] + + result = pd.concat([arr_df, dep_df]) + logger.debug(f"Processed: {len(arr_df)} arrivals, {len(dep_df)} departures") + return result + + +def _recalculate_bus_fields_from_gtfs( + pq_df: pd.DataFrame, service_date: date, local_archive_path: str | None = None +) -> pd.DataFrame: + """Enrich bus LAMP data with GTFS scheduled travel times. + + Replaces LAMP's plan_travel_time_seconds with a scheduled_tt computed from + GTFS (arrival_time - trip start time), matching each LAMP trip to its + nearest scheduled GTFS trip by route/direction/stop. + """ + trip_ids = pq_df["trip_id"].unique() + logger.info(f"Enriching bus LAMP data with GTFS for {len(trip_ids)} unique trips on {service_date}") + gtfs_stops = fetch_stop_times_from_gtfs(trip_ids, service_date, local_archive_path=local_archive_path) + logger.debug(f"Fetched {len(gtfs_stops)} GTFS stop times") + gtfs_stops = gtfs_stops.sort_values(by="arrival_time") + + # Normalize merge-key dtypes so pd.merge_asof's strict by= check doesn't reject string vs. object. + for col in ("route_id", "stop_id", "trip_id"): + pq_df[col] = pq_df[col].astype("string") + gtfs_stops[col] = gtfs_stops[col].astype("string") + pq_df["direction_id"] = pq_df["direction_id"].astype("int16") + gtfs_stops["direction_id"] = gtfs_stops["direction_id"].astype("int16") + + pq_df = pq_df.drop(columns=["scheduled_tt"]) + + trip_start_times = gtfs_stops.groupby("trip_id").arrival_time.transform("min") + gtfs_stops["scheduled_tt"] = gtfs_stops["arrival_time"] - trip_start_times + gtfs_stops["arrival_time"] = gtfs_stops["arrival_time"].astype(float) + + route_starts = pq_df.loc[pq_df.groupby("trip_id").event_time.idxmin()] + route_starts["arrival_time"] = ( + route_starts.event_time - pd.Timestamp(service_date).tz_localize(EASTERN_TIME) + ).dt.total_seconds() + + trip_id_map = pd.merge_asof( + route_starts.sort_values(by="arrival_time"), + gtfs_stops[RTE_DIR_STOP + ["arrival_time", "trip_id"]].drop_duplicates(), + on="arrival_time", + direction="nearest", + by=RTE_DIR_STOP, + suffixes=["", "_scheduled"], + ) + trip_id_map = trip_id_map.drop_duplicates("trip_id").set_index("trip_id").trip_id_scheduled + + pq_df["scheduled_trip_id"] = pq_df.trip_id.map(trip_id_map) + pq_df = pd.merge( + pq_df, + gtfs_stops[["trip_id", "stop_id", "scheduled_tt"]], + how="left", + left_on=["scheduled_trip_id", "stop_id"], + right_on=["trip_id", "stop_id"], + suffixes=["", "_gtfs"], + ) + + missing_tt_mask = pq_df["scheduled_tt"].isna() + if missing_tt_mask.any(): + logger.debug(f"Attempting fallback scheduled_tt matching for {missing_tt_mask.sum()} events") + gtfs_stops["time_bucket"] = (gtfs_stops["arrival_time"] // 1800).astype(int) + bucketed_median_tt = gtfs_stops.groupby(RTE_DIR_STOP + ["time_bucket"])["scheduled_tt"].median() + + event_seconds = ( + pq_df.loc[missing_tt_mask, "event_time"] - pd.Timestamp(service_date).tz_localize(EASTERN_TIME) + ).dt.total_seconds() + event_buckets = (event_seconds // 1800).astype(int) + + lookup_keys = pq_df.loc[missing_tt_mask, RTE_DIR_STOP].copy() + lookup_keys["time_bucket"] = event_buckets.values + fallback_tt = lookup_keys.apply(lambda row: bucketed_median_tt.get(tuple(row)), axis=1) + + pq_df.loc[missing_tt_mask, "scheduled_tt"] = fallback_tt.values + filled_count = (~fallback_tt.isna()).sum() + logger.debug(f"Fallback matching filled {filled_count} events with bucketed median scheduled_tt") + + unmatched_trips = pq_df["scheduled_trip_id"].isna().sum() + if unmatched_trips > 0: + logger.warning(f"{unmatched_trips} bus events could not be matched to a scheduled trip") + return pq_df[BUS_S3_COLUMNS] + + +def _average_bus_scheduled_headways(pq_df: pd.DataFrame, service_date: date) -> pd.DataFrame: + """Bucket scheduled bus headways into 30-minute buckets to smooth the benchmark.""" + logger.debug(f"Calculating average scheduled bus headways for {len(pq_df)} events") + start_time = pd.Timestamp(service_date.year, service_date.month, service_date.day) + end_time = start_time + pd.Timedelta(hours=48) + buckets = pd.date_range(start_time, end_time, freq="30min") + + _enriched = [] + for bucket in buckets: + bucket_start = pd.to_datetime(bucket, unit="s").tz_localize( + EASTERN_TIME, ambiguous=True, nonexistent="shift_forward" + ) + bucket_end = pd.to_datetime(bucket + pd.Timedelta(minutes=30), unit="s").tz_localize( + EASTERN_TIME, ambiguous=True, nonexistent="shift_forward" + ) + filtered = pq_df[(pq_df["event_time"] >= bucket_start) & (pq_df["event_time"] < bucket_end)] + + average_scheduled_headway = filtered.groupby(RTE_DIR_STOP)["scheduled_headway"].mean() + average_scheduled_headway = average_scheduled_headway.round(-1) + + enriched = filtered.merge(average_scheduled_headway, how="left", on=RTE_DIR_STOP, suffixes=["_lamp", ""]) + _enriched.append(enriched) + return pd.concat(_enriched)[BUS_S3_COLUMNS] + + +def ingest_bus_pq_file(df: pd.DataFrame, service_date: date, local_archive_path: str | None = None) -> pd.DataFrame: + """Process and transform columns for a full day's bus events.""" + logger.info(f"Processing {len(df)} raw bus events for service date {service_date}") + + rows_before = len(df) + df = df[df["direction_id"].notna()] + rows_dropped = rows_before - len(df) + if rows_dropped > 0: + logger.warning(f"Dropped {rows_dropped} rows with null direction_id") + + df["direction_id"] = df["direction_id"].astype("int16") + df["service_date"] = df["service_date"].astype(str) + df = df.rename(columns=BUS_COLUMN_RENAME_MAP) + + logger.info("Processing arrival/departure times") + processed = _process_bus_arrival_departure_times(df) + events_before = len(processed) + processed = processed[processed["stop_id"].notna()] + events_dropped = events_before - len(processed) + if events_dropped > 0: + logger.warning(f"Dropped {events_dropped} events with null stop_id") + + logger.info("Recalculating fields from GTFS") + processed = _recalculate_bus_fields_from_gtfs(processed, service_date, local_archive_path) + + logger.info("Averaging scheduled headways") + processed = _average_bus_scheduled_headways(processed, service_date) + + logger.info(f"Processing complete: {len(processed)} events ready for upload") + return processed.sort_values(by=["event_time"]) + + +def upload_bus_to_s3(group_key_and_events: Tuple[tuple, pd.DataFrame], service_date: date) -> None: + """Upload bus events to S3, grouped by route-direction-stop.""" + (route_id, direction_id, stop_id), stop_events = group_key_and_events + + s3_key = BUS_S3_KEY_TEMPLATE.format( + route_id=route_id, + direction_id=direction_id, + stop_id=stop_id, + YYYY=service_date.year, + _M=service_date.month, + _D=service_date.day, + ) + logger.debug(f"Uploading {len(stop_events)} events for {route_id}-{direction_id}-{stop_id}") + try: + s3.upload_df_as_csv(BUS_S3_BUCKET, s3_key, stop_events) + except Exception as e: + logger.error(f"Failed to upload bus events for {route_id}-{direction_id}-{stop_id}: {e}") + raise + return [(route_id, direction_id, stop_id)] + + +_parallel_upload = parallel.make_parallel(upload_bus_to_s3) + + +def ingest_bus_data(service_date: date, local_archive_path: str | None = None): + """Ingest and upload bus LAMP data for a given service date.""" + logger.info(f"Starting bus LAMP data ingestion for service date {service_date}") + try: + df = fetch_bus_pq_file_from_remote(service_date) + except ValueError as e: + logger.error(f"Failed to fetch bus data for {service_date}: {e}") + return + except Exception as e: + logger.exception(f"Unexpected error fetching bus data for {service_date}: {e}") + raise + + try: + processed = ingest_bus_pq_file(df, service_date, local_archive_path=local_archive_path) + except Exception as e: + logger.exception(f"Error processing bus data for {service_date}: {e}") + raise + + # Group by route-direction-stop and parallel upload to S3 + group_event_groups = processed.groupby(RTE_DIR_STOP) + num_groups = len(group_event_groups) + logger.info(f"Uploading events for {num_groups} route-direction-stop groups to S3") + try: + _parallel_upload(group_event_groups, service_date) + except Exception as e: + logger.exception(f"Error uploading bus data for {service_date}: {e}") + raise + logger.info(f"Bus LAMP data ingestion complete for service date {service_date}") + + +def ingest_today_bus_data(): + """Ingest and upload today's bus LAMP data.""" + service_date = get_current_service_date() + logger.info(f"Ingesting today's bus data (service date: {service_date})") + ingest_bus_data(service_date) + + +def ingest_yesterday_bus_data(): + """Ingest and upload yesterday's bus LAMP data.""" + service_date = get_current_service_date() - pd.Timedelta(days=1) + logger.info(f"Ingesting yesterday's bus data (service date: {service_date})") + ingest_bus_data(service_date) + + +if __name__ == "__main__": + import os + + log_level = os.environ.get("LOG_LEVEL", "INFO").upper() + logging.basicConfig( + level=getattr(logging, log_level, logging.INFO), + format="%(asctime)s - %(levelname)s: %(message)s", + datefmt="%H:%M:%S", + ) + ingest_today_bus_data() diff --git a/mbta-performance/chalicelib/lamp/tests/sample_data/bus-20260407-sample.parquet b/mbta-performance/chalicelib/lamp/tests/sample_data/bus-20260407-sample.parquet new file mode 100644 index 0000000..5a9de64 Binary files /dev/null and b/mbta-performance/chalicelib/lamp/tests/sample_data/bus-20260407-sample.parquet differ diff --git a/mbta-performance/chalicelib/lamp/tests/test_bus_ingest.py b/mbta-performance/chalicelib/lamp/tests/test_bus_ingest.py new file mode 100644 index 0000000..e4d8e84 --- /dev/null +++ b/mbta-performance/chalicelib/lamp/tests/test_bus_ingest.py @@ -0,0 +1,165 @@ +import io +import os +import unittest +from datetime import date +from unittest import mock + +import pandas as pd + +from .. import bus_constants, bus_ingest + +DATA_PREFIX = os.path.join(os.path.dirname(__file__), "sample_data") +SAMPLE_BUS_DATA_PATH = os.path.join(DATA_PREFIX, "bus-20260407-sample.parquet") + + +def _empty_gtfs_mock() -> pd.DataFrame: + """Empty GTFS stop_times mock matching fetch_stop_times_from_gtfs's schema.""" + return pd.DataFrame( + { + "trip_id": pd.array([], dtype="string"), + "stop_id": pd.array([], dtype="string"), + "arrival_time": pd.array([], dtype="Int64"), + "route_id": pd.array([], dtype="string"), + "direction_id": pd.array([], dtype="int16"), + } + ) + + +class TestBusIngest(unittest.TestCase): + def setUp(self): + with open(SAMPLE_BUS_DATA_PATH, "rb") as f: + self.data = f.read() + + self.sample_df = pd.read_parquet( + io.BytesIO(self.data), + columns=bus_constants.BUS_LAMP_COLUMNS, + engine="pyarrow", + dtype_backend="numpy_nullable", + ) + self.mock_gtfs_data = _empty_gtfs_mock() + + def test_fetch_bus_pq_file_from_remote(self): + mock_response = mock.Mock(status_code=200, content=self.data) + with mock.patch("requests.get", return_value=mock_response): + df = bus_ingest.fetch_bus_pq_file_from_remote(date(2026, 4, 7)) + self.assertEqual(set(df.columns), set(bus_constants.BUS_LAMP_COLUMNS)) + self.assertGreater(len(df), 0) + + def test_fetch_bus_pq_file_from_remote_failure(self): + mock_response = mock.Mock(status_code=404) + with mock.patch("requests.get", return_value=mock_response): + with self.assertRaises(ValueError) as context: + bus_ingest.fetch_bus_pq_file_from_remote(date(2026, 4, 7)) + self.assertIn("Failed to fetch bus LAMP parquet file", str(context.exception)) + + def test_process_bus_arrival_departure_times(self): + df = self.sample_df.rename(columns=bus_constants.BUS_COLUMN_RENAME_MAP) + result = bus_ingest._process_bus_arrival_departure_times(df) + + arrivals = result[result.event_type == "ARR"] + departures = result[result.event_type == "DEP"] + + self.assertGreater(len(arrivals), 0) + self.assertGreater(len(departures), 0) + self.assertListEqual(list(result.columns), bus_constants.BUS_S3_COLUMNS) + + def test_process_bus_arrival_departure_times_timezone(self): + df = self.sample_df.rename(columns=bus_constants.BUS_COLUMN_RENAME_MAP) + result = bus_ingest._process_bus_arrival_departure_times(df) + + # All event_times should be in Eastern Time + for event_time in result["event_time"].dropna().head(5): + self.assertEqual(str(event_time.tzinfo), "US/Eastern") + + def test_departures_use_previous_stop_id(self): + df = self.sample_df.rename(columns=bus_constants.BUS_COLUMN_RENAME_MAP) + # Find rows where previous_stop_id differs from stop_id + has_prev = df[df["previous_stop_id"].notna() & (df["previous_stop_id"] != df["stop_id"])] + if len(has_prev) > 0: + result = bus_ingest._process_bus_arrival_departure_times(df) + departures = result[result.event_type == "DEP"] + # DEP events should use previous_stop_id, not the original stop_id + self.assertGreater(len(departures), 0) + + def test_ingest_bus_pq_file(self): + with mock.patch("chalicelib.lamp.bus_ingest.fetch_stop_times_from_gtfs", return_value=self.mock_gtfs_data): + result = bus_ingest.ingest_bus_pq_file(self.sample_df, date(2026, 4, 7)) + + # No null stop_ids + self.assertFalse(result["stop_id"].isna().any()) + # Output columns match + self.assertListEqual(list(result.columns), bus_constants.BUS_S3_COLUMNS) + # Sorted by event_time + event_times = result["event_time"].tolist() + self.assertEqual(event_times, sorted(event_times)) + # service_date is a string + for sdate in result["service_date"].unique(): + self.assertIsInstance(sdate, str) + + def test_upload_bus_to_s3_key_format(self): + df = pd.DataFrame({col: ["test"] for col in bus_constants.BUS_S3_COLUMNS}) + + with mock.patch("chalicelib.lamp.bus_ingest.s3.upload_df_as_csv") as mock_upload: + result = bus_ingest.upload_bus_to_s3((("1", 0, "110"), df), date(2026, 4, 7)) + + mock_upload.assert_called_once() + call_args = mock_upload.call_args + self.assertEqual(call_args[0][0], "tm-mbta-performance") + expected_key = "Events-lamp/bus-daily-data/1-0-110/Year=2026/Month=4/Day=7/events.csv" + self.assertEqual(call_args[0][1], expected_key) + self.assertEqual(result, [("1", 0, "110")]) + + def test_ingest_bus_data_end_to_end(self): + mock_response = mock.Mock(status_code=200, content=self.data) + with mock.patch("requests.get", return_value=mock_response): + with mock.patch("chalicelib.lamp.bus_ingest._parallel_upload") as mock_upload: + with mock.patch( + "chalicelib.lamp.bus_ingest.fetch_stop_times_from_gtfs", + return_value=self.mock_gtfs_data, + ): + bus_ingest.ingest_bus_data(date(2026, 4, 7)) + mock_upload.assert_called_once() + + def test_ingest_bus_data_no_file_found(self): + mock_response = mock.Mock(status_code=404) + with mock.patch("requests.get", return_value=mock_response): + # Should not raise - logs error and returns + bus_ingest.ingest_bus_data(date(2026, 4, 7)) + + def test_ingest_today_bus_data(self): + mock_response = mock.Mock(status_code=200, content=self.data) + with mock.patch("requests.get", return_value=mock_response): + with mock.patch("chalicelib.lamp.bus_ingest._parallel_upload"): + with mock.patch( + "chalicelib.lamp.bus_ingest.fetch_stop_times_from_gtfs", return_value=self.mock_gtfs_data + ): + with mock.patch( + "chalicelib.lamp.bus_ingest.get_current_service_date", return_value=date(2026, 4, 7) + ): + bus_ingest.ingest_today_bus_data() + + def test_ingest_yesterday_bus_data(self): + mock_response = mock.Mock(status_code=200, content=self.data) + with mock.patch("requests.get", return_value=mock_response): + with mock.patch("chalicelib.lamp.bus_ingest._parallel_upload"): + with mock.patch( + "chalicelib.lamp.bus_ingest.fetch_stop_times_from_gtfs", return_value=self.mock_gtfs_data + ): + with mock.patch( + "chalicelib.lamp.bus_ingest.get_current_service_date", return_value=date(2026, 4, 8) + ): + bus_ingest.ingest_yesterday_bus_data() + + def test_column_rename_map(self): + df = self.sample_df.copy() + # Verify source columns exist + self.assertIn("stopped_duration_seconds", df.columns) + self.assertIn("route_direction_headway_seconds", df.columns) + self.assertIn("plan_travel_time_seconds", df.columns) + self.assertIn("plan_route_direction_headway_seconds", df.columns) + + renamed = df.rename(columns=bus_constants.BUS_COLUMN_RENAME_MAP) + self.assertIn("dwell_time_seconds", renamed.columns) + self.assertIn("headway_seconds", renamed.columns) + self.assertIn("scheduled_tt", renamed.columns) + self.assertIn("scheduled_headway", renamed.columns)