diff --git a/config/cron.sh b/config/cron.sh index 955178b3..81a0f73a 100644 --- a/config/cron.sh +++ b/config/cron.sh @@ -1,5 +1,21 @@ -# This script adds commands to the current crontab +#!/bin/sh -# run the daily script at 1am every morning -# TODO: make sure timezone is PST -crontab -l | { cat; echo "0 1 * * * /home/csss-site/csss-site-backend/src/cron/daily.py"; } | crontab - +# Install the weekly TransLink static schedule refresh without duplicating it. +set -eu + +cron_file=$(mktemp) +trap 'rm -f "$cron_file"' EXIT + +crontab -l 2>/dev/null | sed \ + -e '/# BEGIN CSSS TRANSLINK STATIC/,/# END CSSS TRANSLINK STATIC/d' \ + -e '\|scripts.refresh_translink_static|d' > "$cron_file" || true + +{ + cat "$cron_file" + printf '%s\n' \ + '# BEGIN CSSS TRANSLINK STATIC' \ + 'PATH=/home/csss-site/.local/bin:/usr/local/bin:/usr/bin:/bin' \ + 'CRON_TZ=America/Vancouver' \ + '0 23 * * 5 cd /home/csss-site/csss-site-backend/src && uv run python -m scripts.refresh_translink_static' \ + '# END CSSS TRANSLINK STATIC' +} | crontab - diff --git a/src/scripts/refresh_translink_static.py b/src/scripts/refresh_translink_static.py new file mode 100644 index 00000000..70b61d93 --- /dev/null +++ b/src/scripts/refresh_translink_static.py @@ -0,0 +1,47 @@ +"""Download, preprocess, and store the TransLink static GTFS schedule.""" + +import asyncio +import logging + +import httpx + +import database +from translink.crud import refresh_static_schedule + +_logger = logging.getLogger(__name__) + + +async def refresh() -> None: + await database.setup_database() + if database.sessionmanager is None: + raise RuntimeError("Database has not been initialized") + + manager = database.sessionmanager + try: + async with httpx.AsyncClient(timeout=60, follow_redirects=True) as client: + async with manager.session() as session: + cache = await refresh_static_schedule(session, client) + departure_count = sum(len(rows) for rows in cache["departures"].values()) + _logger.info( + "Stored TransLink static schedule version %s with %s departures covering %s through %s", + cache["version"], + departure_count, + cache["coverage"]["start_date"], + cache["coverage"]["end_date"], + ) + finally: + await manager.close() + + +def main() -> int: + logging.basicConfig(level=logging.INFO) + try: + asyncio.run(refresh()) + except Exception: + _logger.exception("Failed to refresh the TransLink static schedule") + return 1 + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/src/translink/README.md b/src/translink/README.md index 43d76eb7..c9bcb83e 100644 --- a/src/translink/README.md +++ b/src/translink/README.md @@ -9,8 +9,30 @@ All dates are adjusted for the America/Vancouver timezone. 3. Make sure your database has the correct migrations `alembic upgrade head`. Reload your test database as well `python src/load_test_db.py` 4. Start (or restart) the web server to test the endpoints +## Static schedule refresh + +The static GTFS archive is downloaded and preprocessed outside HTTP requests. Before serving the static schedule for +the first time, populate the cache manually from the `src` directory: + +```bash +# in ./src +uv run python -m scripts.refresh_translink_static +``` + +If deploying on the web server run the cron job. +```bash +# in root +sh config/cron.sh +``` +Production refreshes it every Friday at 11:00 PM in the America/Vancouver timezone. Install or update that cron entry +by running `sh config/cron.sh` from the repository root. The installer is idempotent. + +If a refresh fails, the prior database row is preserved and the command exits unsuccessfully. If no compatible cache +can serve the current date, the static and combined schedule endpoints return HTTP 503; requests never download or +parse the static GTFS archive. + ## Endpoints You can see the exact schemas in the `/docs` page. At the time this was written there are three endpoints: 1. `translink/realtime`: returns realtime data for buses that are at or are approaching SFU -2. `translink/static`: returns the schedule for the current day +2. `translink/static`: returns the preprocessed schedule for the current day 3. `translink/schedule`: combines the realtime and static data to show if a bus is at the loop, is running late, or was cancelled diff --git a/src/translink/crud.py b/src/translink/crud.py index ee9cce60..43ac5cd5 100644 --- a/src/translink/crud.py +++ b/src/translink/crud.py @@ -25,6 +25,18 @@ REALTIME_CACHE_ID = 1 REALTIME_CACHE_TTL_SECONDS = 90 REALTIME_CACHE_LOCK_ID = 2026062601 +STATIC_CACHE_ID = 1 +STATIC_CACHE_VERSION = 1 +STATIC_CACHE_UNAVAILABLE_MESSAGE = "static TransLink schedule cache is unavailable" + +WEEKDAYS = ("monday", "tuesday", "wednesday", "thursday", "friday", "saturday", "sunday") + +type StaticScheduleEntry = dict[str, str | int] +type StaticScheduleCache = dict[str, Any] + + +class StaticScheduleCacheUnavailableError(RuntimeError): + """Raised when the preprocessed static schedule cannot serve a date.""" # Taken from the static data. @@ -49,131 +61,228 @@ def _gtfs_time_to_seconds(time_str: str) -> int: return h * 3600 + m * 60 + s -def _get_active_service_ids(z: zipfile.ZipFile) -> set[str]: - today = datetime.now(tz=TZ_INFO) - # Dates in the calendar.txt are in YYYYMMDD - date_str = today.strftime("%Y%m%d") - day_name = today.strftime("%A").lower() - - calendar = pd.read_csv(z.open("calendar.txt"), dtype=str) - active = set( - calendar[ - (calendar[day_name] == "1") & (calendar["start_date"] <= date_str) & (calendar["end_date"] >= date_str) - ]["service_id"] - ) +def parse_static_schedule(content: bytes) -> StaticScheduleCache: + """Reduce a GTFS archive to the service rules and departures used by the kiosk.""" + try: + with zipfile.ZipFile(io.BytesIO(content)) as archive: + filenames = set(archive.namelist()) + if "calendar.txt" not in filenames and "calendar_dates.txt" not in filenames: + raise ValueError("GTFS archive contains neither calendar.txt nor calendar_dates.txt") + + calendar = ( + pd.read_csv( + archive.open("calendar.txt"), + dtype=str, + usecols=["service_id", "start_date", "end_date", *WEEKDAYS], + ) + if "calendar.txt" in filenames + else pd.DataFrame(columns=["service_id", "start_date", "end_date", *WEEKDAYS]) + ) + exceptions = ( + pd.read_csv( + archive.open("calendar_dates.txt"), + dtype=str, + usecols=["service_id", "date", "exception_type"], + ) + if "calendar_dates.txt" in filenames + else pd.DataFrame(columns=["service_id", "date", "exception_type"]) + ) + trips = pd.read_csv( + archive.open("trips.txt"), + dtype=str, + usecols=["trip_id", "route_id", "service_id", "direction_id"], + ) + stop_times = pd.read_csv( + archive.open("stop_times.txt"), + dtype=str, + usecols=["trip_id", "stop_id", "departure_time"], + ) + except (KeyError, ValueError, zipfile.BadZipFile, pd.errors.ParserError) as e: + raise RuntimeError(f"Failed to parse static schedule: {e}") from e - # These are exceptions to services in the calendar - # exception_type=1 means service was added - # exception_type=2 means service was removed - exceptions = pd.read_csv(z.open("calendar_dates.txt"), dtype=str) - added = exceptions[(exceptions["date"] == date_str) & (exceptions["exception_type"] == "1")]["service_id"] - removed = exceptions[(exceptions["date"] == date_str) & (exceptions["exception_type"] == "2")]["service_id"] - active |= set(added) - active -= set(removed) - return active + route_ids = set(BUS_DATA) + filtered_trips = trips[trips["route_id"].isin(route_ids)].copy() + expected_directions = filtered_trips["route_id"].map(lambda route_id: str(BUS_DATA[route_id][0])) + filtered_trips = filtered_trips[filtered_trips["direction_id"] == expected_directions] + stop_ids = {data[1] for data in BUS_DATA.values()} + relevant_stop_times = stop_times[ + stop_times["trip_id"].isin(filtered_trips["trip_id"]) & stop_times["stop_id"].isin(stop_ids) + ] + merged = relevant_stop_times.merge(filtered_trips[["trip_id", "route_id", "service_id"]], on="trip_id") + expected_stops = merged["route_id"].map(lambda route_id: BUS_DATA[route_id][1]) + merged = merged[merged["stop_id"] == expected_stops].copy() -async def fetch_static_schedule(client: AsyncClient) -> pd.DataFrame: - """ - Gets the static bus schedule from the static TransLink GTFS API - """ - # Retrieve the static TransLink bus schedule data - try: - static_response = await client.get(STATIC_URL) - except httpx.HTTPError as e: - raise RuntimeError(f"Failed to fetch static schedule: {e}") from e + if merged.empty: + raise RuntimeError("Static schedule contains no departures for the configured routes") + merged["bus_number"] = merged["route_id"].map(lambda route_id: BUS_DATA[route_id][2]) try: - z = zipfile.ZipFile(io.BytesIO(static_response.content)) - except zipfile.BadZipFile as e: - raise RuntimeError(f"Failed to read static schedule zip file: {e}") from e - - # A trip is from one stop to the next one - active_services = _get_active_service_ids(z) - trips = pd.read_csv(z.open("trips.txt"), dtype=str) - # Stop times contain when the bus should depart a bus stop - stop_times = pd.read_csv(z.open("stop_times.txt"), dtype=str) - - # From all the active trips, only get the ones that go to the bus loop - route_ids = set(BUS_DATA.keys()) - filtered_trips = trips[trips["route_id"].isin(list(route_ids)) & trips["service_id"].isin(list(active_services))] - filtered_trips = filtered_trips[ - filtered_trips.apply(lambda row: int(row["direction_id"]) == BUS_DATA[row["route_id"]][0], axis=1) + merged["departure_seconds"] = merged["departure_time"].map(_gtfs_time_to_seconds) + except (AttributeError, TypeError, ValueError) as e: + raise RuntimeError(f"Failed to parse static schedule departure times: {e}") from e + + relevant_service_ids = set(merged["service_id"]) + calendar = calendar[calendar["service_id"].isin(relevant_service_ids)] + exceptions = exceptions[exceptions["service_id"].isin(relevant_service_ids)] + + services: dict[str, dict[str, Any]] = {} + for row in calendar.to_dict(orient="records"): + services[row["service_id"]] = { + "start_date": row["start_date"], + "end_date": row["end_date"], + "weekdays": [index for index, weekday in enumerate(WEEKDAYS) if row[weekday] == "1"], + } + + exception_map: dict[str, dict[str, list[str]]] = {} + for row in exceptions.to_dict(orient="records"): + exception = exception_map.setdefault(row["date"], {"added": [], "removed": []}) + if row["exception_type"] == "1": + exception["added"].append(row["service_id"]) + elif row["exception_type"] == "2": + exception["removed"].append(row["service_id"]) + + departure_columns = ["trip_id", "route_id", "bus_number", "departure_time", "departure_seconds"] + departures: dict[str, list[StaticScheduleEntry]] = {} + for service_id, rows in merged.groupby("service_id"): + schedule = cast(pd.DataFrame, rows[departure_columns]).sort_values(by=["route_id", "departure_seconds"]) + departures[str(service_id)] = cast(list[StaticScheduleEntry], schedule.to_dict(orient="records")) + + coverage_dates = [ + *(str(value) for value in calendar["start_date"]), + *(str(value) for value in calendar["end_date"]), + *(str(value) for value in exceptions["date"]), ] + if not coverage_dates: + raise RuntimeError("Static schedule contains no calendar coverage for the configured routes") - # Get the stop times entries for the stops at the bus loop - stop_ids = {s[1] for s in BUS_DATA.values()} - stop_times = stop_times[ - stop_times["trip_id"].isin(list(filtered_trips["trip_id"])) & stop_times["stop_id"].isin(list(stop_ids)) - ] - - # Join the data from the trips and the stops - # Casts are done to avoid some typing issues, but they might be unnecessary - merged = stop_times.merge(cast(pd.DataFrame, filtered_trips[["trip_id", "route_id"]]), on="trip_id") - merged = cast( - pd.DataFrame, merged[merged.apply(lambda row: row["stop_id"] == BUS_DATA[row["route_id"]][1], axis=1)] - ) # filter for the stops we care about - - merged = merged.copy() # stops pandas from complaining about modifying original data - merged["bus_number"] = merged["route_id"].map(lambda r: BUS_DATA[r][2]) - merged["departure_seconds"] = merged["departure_time"].map(_gtfs_time_to_seconds) - return ( - cast(pd.DataFrame, merged[["trip_id", "route_id", "bus_number", "departure_time", "departure_seconds"]]) - .reset_index(drop=True) - .sort_values(by=["route_id", "departure_seconds"]) - ) + return { + "version": STATIC_CACHE_VERSION, + "coverage": {"start_date": min(coverage_dates), "end_date": max(coverage_dates)}, + "services": services, + "exceptions": exception_map, + "departures": departures, + } -async def get_or_fetch_static_schedule(db_session: DBSession, client: AsyncClient) -> tuple[date, pd.DataFrame]: - today = datetime.now(tz=TZ_INFO).date() - +async def fetch_static_schedule(client: AsyncClient) -> StaticScheduleCache: + """Download and preprocess the static TransLink GTFS feed.""" try: - result = await db_session.scalar( - sqlalchemy.select(TransLinkStaticScheduleDB).where(TransLinkStaticScheduleDB.date_fetched == today) - ) - except sqlalchemy.exc.SQLAlchemyError as e: - logging.error(f"Failed to query static schedule from database: {e}") - result = None + response = await client.get(STATIC_URL) + response.raise_for_status() + except httpx.HTTPError as e: + raise RuntimeError(f"Failed to fetch static schedule: {e}") from e - if result is not None: - return (result.date_fetched, pd.DataFrame(result.schedule)) + return parse_static_schedule(response.content) - result = await fetch_static_schedule(client) - if result.empty: - raise ValueError("No active schedule found for today") +async def refresh_static_schedule(db_session: DBSession, client: AsyncClient) -> StaticScheduleCache: + """Fetch and atomically replace the preprocessed static schedule cache.""" + schedule = await fetch_static_schedule(client) try: await db_session.merge( - TransLinkStaticScheduleDB(id=1, date_fetched=today, schedule=result.to_dict(orient="records")) + TransLinkStaticScheduleDB( + id=STATIC_CACHE_ID, + date_fetched=datetime.now(tz=TZ_INFO).date(), + schedule=schedule, + ) ) await db_session.commit() except sqlalchemy.exc.SQLAlchemyError as e: - logging.warning(f"Failed to cache static schedule to database: {e}") await db_session.rollback() + raise RuntimeError(f"Failed to store static schedule: {e}") from e + return schedule + + +def resolve_static_schedule(cache: StaticScheduleCache, service_date: date) -> list[StaticScheduleEntry]: + """Resolve a preprocessed weekly cache into departures for one service date.""" + try: + if cache["version"] != STATIC_CACHE_VERSION: + raise ValueError(f"unsupported cache version {cache['version']}") + + date_str = service_date.strftime("%Y%m%d") + coverage = cache["coverage"] + if not isinstance(coverage, dict) or not all( + isinstance(coverage.get(key), str) and len(coverage[key]) == 8 for key in ("start_date", "end_date") + ): + raise ValueError("invalid cache coverage") + if not coverage["start_date"] <= date_str <= coverage["end_date"]: + raise ValueError(f"date {date_str} is outside cache coverage") + + services = cache["services"] + departures = cache["departures"] + if not isinstance(services, dict) or not isinstance(departures, dict): + raise ValueError("invalid cache services or departures") + active_services = { + service_id + for service_id, service in services.items() + if service["start_date"] <= date_str <= service["end_date"] + and service_date.weekday() in service["weekdays"] + } + exceptions = cache["exceptions"].get(date_str, {"added": [], "removed": []}) + if not isinstance(exceptions["added"], list) or not isinstance(exceptions["removed"], list): + raise ValueError("invalid cache exceptions") + active_services.update(exceptions["added"]) + active_services.difference_update(exceptions["removed"]) + + schedule = [row for service_id in active_services for row in departures.get(service_id, [])] + required_string_fields = ("trip_id", "route_id", "bus_number", "departure_time") + if any( + not isinstance(row, dict) + or any(not isinstance(row.get(field), str) for field in required_string_fields) + or not isinstance(row.get("departure_seconds"), int) + for row in schedule + ): + raise ValueError("invalid cached departure") + return sorted(schedule, key=lambda row: (str(row["route_id"]), int(row["departure_seconds"]))) + except (AttributeError, KeyError, TypeError, ValueError) as e: + raise StaticScheduleCacheUnavailableError(STATIC_CACHE_UNAVAILABLE_MESSAGE) from e + + +async def get_static_schedule( + db_session: DBSession, service_date: date | None = None +) -> tuple[date, list[StaticScheduleEntry]]: + """Read the weekly cache and resolve it for a date without network or pandas work.""" + target_date = service_date or datetime.now(tz=TZ_INFO).date() + try: + cached = await db_session.scalar( + sqlalchemy.select(TransLinkStaticScheduleDB).where(TransLinkStaticScheduleDB.id == STATIC_CACHE_ID) + ) + except sqlalchemy.exc.SQLAlchemyError as e: + logging.error("Failed to query static schedule cache: %s", e) + raise StaticScheduleCacheUnavailableError(STATIC_CACHE_UNAVAILABLE_MESSAGE) from e - return (today, result) + if cached is None: + raise StaticScheduleCacheUnavailableError(STATIC_CACHE_UNAVAILABLE_MESSAGE) + return target_date, resolve_static_schedule(cached.schedule, target_date) -def get_next_departures(schedule: pd.DataFrame, n: int = 3) -> pd.DataFrame: +def get_next_departures(schedule: list[StaticScheduleEntry], n: int = 3) -> list[StaticScheduleEntry]: """ Get the next few departures for today. Args: - schedule: static schedule filtered out for the relevant routes + schedule: static schedule filtered for the relevant routes n: the number of departures to get for each route Returns: - A dataframe with the next n departures for each route, sorted by route ID and departure time (in seconds). + The next n departures for each route, sorted by route ID and departure time (in seconds). """ now = datetime.now(tz=TZ_INFO) current_seconds = int((now - now.replace(hour=0, minute=0, second=0, microsecond=0)).total_seconds()) - upcoming = cast(pd.DataFrame, schedule[schedule["departure_seconds"] > current_seconds]) - return ( - upcoming.sort_values("departure_seconds") - .groupby("route_id") - .head(n) - .sort_values(["route_id", "departure_seconds"]) + upcoming = sorted( + (row for row in schedule if int(row["departure_seconds"]) > current_seconds), + key=lambda row: int(row["departure_seconds"]), ) + route_counts: dict[str, int] = {} + result: list[StaticScheduleEntry] = [] + for row in upcoming: + route_id = str(row["route_id"]) + if route_counts.get(route_id, 0) >= n: + continue + result.append(row) + route_counts[route_id] = route_counts.get(route_id, 0) + 1 + return sorted(result, key=lambda row: (str(row["route_id"]), int(row["departure_seconds"]))) def _parse_feed(content: bytes) -> FeedMessage: @@ -313,12 +422,12 @@ def _response_from_static_row(row: Any, delay: int = 0, status: BusStatus = BusS status=status, ) - _, schedule = await get_or_fetch_static_schedule(db_session, client) + _, schedule = await get_static_schedule(db_session) next_departures = get_next_departures(schedule) trip_feed = await get_or_fetch_realtime_feed(db_session, client) # If the trip feed fails to fetch then just return information from the static schedule. if trip_feed is None: - return [_response_from_static_row(row) for _, row in next_departures.iterrows()] + return [_response_from_static_row(row) for row in next_departures] # FeedMessage is generated at runtime, so the type checker can't find this function # Map all the realtime data to each bus's status @@ -357,5 +466,5 @@ def _response_from_static_row(row: Any, delay: int = 0, status: BusStatus = BusS row, *realtime_map.get(cast(str, row["trip_id"]), (0, BusStatus.OnTime)), ) - for _, row in next_departures.iterrows() + for row in next_departures ] diff --git a/src/translink/tables.py b/src/translink/tables.py index b00b777b..cbb370ee 100644 --- a/src/translink/tables.py +++ b/src/translink/tables.py @@ -1,4 +1,5 @@ from datetime import date, datetime +from typing import Any from sqlalchemy import DateTime, LargeBinary from sqlalchemy.dialects.postgresql import JSONB @@ -13,7 +14,7 @@ class TransLinkStaticScheduleDB(Base): id: Mapped[int] = mapped_column(primary_key=True) date_fetched: Mapped[date] = mapped_column() - schedule: Mapped[list[dict]] = mapped_column(JSONB) + schedule: Mapped[dict[str, Any]] = mapped_column(JSONB) class TransLinkRealtimeCacheDB(Base): diff --git a/src/translink/urls.py b/src/translink/urls.py index 2fa2d566..4a626ddf 100644 --- a/src/translink/urls.py +++ b/src/translink/urls.py @@ -1,10 +1,12 @@ -from fastapi import APIRouter, Request +from fastapi import APIRouter, HTTPException, Request, status from database import DBSession from translink.crud import ( + STATIC_CACHE_UNAVAILABLE_MESSAGE, + StaticScheduleCacheUnavailableError, fetch_realtime_schedule, get_departure_statuses, - get_or_fetch_static_schedule, + get_static_schedule, ) from translink.models import ( TransLinkRealtimeResponse, @@ -37,19 +39,31 @@ async def get_realtime_schedule(db_session: DBSession, request: Request): response_model=TransLinkStaticResponse, operation_id="get_static_schedule", ) -async def get_static_schedule(db_session: DBSession, request: Request): - date_fetched, df = await get_or_fetch_static_schedule(db_session, request.app.state.http_client) - schedule = [TransLinkStaticScheduleEntry(**row) for row in df.to_dict(orient="records")] +async def get_static_schedule_endpoint(db_session: DBSession): + try: + date_fetched, rows = await get_static_schedule(db_session) + except StaticScheduleCacheUnavailableError as e: + raise HTTPException( + status_code=status.HTTP_503_SERVICE_UNAVAILABLE, + detail=STATIC_CACHE_UNAVAILABLE_MESSAGE, + ) from e + schedule = [TransLinkStaticScheduleEntry(**row) for row in rows] return TransLinkStaticResponse(date_fetched=date_fetched, schedule=schedule) @router.get( "/schedule", - description="Get the departure schedule with bus status. Attempts to use the cached static schedule first.", + description="Get the departure schedule with bus status using the preprocessed static schedule cache.", response_description="The next three depature times with bus status information.", response_model=list[TransLinkScheduleResponse], operation_id="get_departure_schedule", ) async def get_departure_schedule(db_session: DBSession, request: Request): - return await get_departure_statuses(db_session, request.app.state.http_client) + try: + return await get_departure_statuses(db_session, request.app.state.http_client) + except StaticScheduleCacheUnavailableError as e: + raise HTTPException( + status_code=status.HTTP_503_SERVICE_UNAVAILABLE, + detail=STATIC_CACHE_UNAVAILABLE_MESSAGE, + ) from e diff --git a/tests/unit/test_translink.py b/tests/unit/test_translink.py index de5fa350..efb426eb 100644 --- a/tests/unit/test_translink.py +++ b/tests/unit/test_translink.py @@ -1,6 +1,6 @@ import io import zipfile -from datetime import datetime, timedelta +from datetime import date, datetime, timedelta from unittest.mock import AsyncMock, MagicMock, patch import httpx @@ -13,13 +13,18 @@ from constants import TZ_INFO from translink.crud import ( BUS_DATA, + STATIC_CACHE_UNAVAILABLE_MESSAGE, + STATIC_CACHE_VERSION, + StaticScheduleCacheUnavailableError, _gtfs_time_to_seconds, fetch_realtime_schedule, fetch_static_schedule, get_departure_statuses, get_next_departures, get_or_fetch_realtime_feed, - get_or_fetch_static_schedule, + get_static_schedule, + refresh_static_schedule, + resolve_static_schedule, ) from translink.models import BusStatus, TransLinkRealtimeResponse, TransLinkScheduleResponse from translink.tables import TransLinkRealtimeCacheDB, TransLinkStaticScheduleDB @@ -36,7 +41,7 @@ def _current_day_name() -> str: return datetime.now(tz=TZ_INFO).strftime("%A").lower() -def make_gtfs_zip(departure_time: str = "23:00:00") -> bytes: +def make_gtfs_zip(departure_time: str = "23:00:00", active_weekdays: set[int] | None = None) -> bytes: """ Return a minimal but valid GTFS zip whose single service is active today, with one trip per route in BUS_DATA. @@ -46,12 +51,11 @@ def make_gtfs_zip(departure_time: str = "23:00:00") -> bytes: includes them. """ buf = io.BytesIO() - day = _current_day_name() all_days = ["monday", "tuesday", "wednesday", "thursday", "friday", "saturday", "sunday"] + active_weekdays = active_weekdays or {datetime.now(tz=TZ_INFO).weekday()} with zipfile.ZipFile(buf, "w") as z: - # calendar.txt - one service active only on today's weekday - cal_row = {d: ("1" if d == day else "0") for d in all_days} + cal_row = {day: ("1" if index in active_weekdays else "0") for index, day in enumerate(all_days)} cal_row.update({"service_id": "SVC1", "start_date": "20240101", "end_date": "20991231"}) z.writestr("calendar.txt", pd.DataFrame([cal_row]).to_csv(index=False)) @@ -136,6 +140,29 @@ def mock_db_session(cached_row=None) -> AsyncMock: return session +def make_static_cache( + schedule: list[dict], + service_date: date | None = None, + *, + version: int = STATIC_CACHE_VERSION, +) -> dict: + target_date = service_date or datetime.now(tz=TZ_INFO).date() + date_str = target_date.strftime("%Y%m%d") + return { + "version": version, + "coverage": {"start_date": date_str, "end_date": date_str}, + "services": { + "SVC1": { + "start_date": date_str, + "end_date": date_str, + "weekdays": [target_date.weekday()], + } + }, + "exceptions": {}, + "departures": {"SVC1": schedule}, + } + + # --------------------------------------------------------------------------- # Unit tests — pure functions # --------------------------------------------------------------------------- @@ -160,30 +187,28 @@ async def test__get_next_departures_filters_past(): midnight = now.replace(hour=0, minute=0, second=0, microsecond=0) now_secs = int((now - midnight).total_seconds()) - schedule = pd.DataFrame( - [ - # Already departed - must be excluded - { - "trip_id": "past_trip", - "route_id": "6656", - "bus_number": "143", - "departure_time": "00:01:00", - "departure_seconds": 60, - }, - # Future - must be included - { - "trip_id": "future_trip", - "route_id": "6656", - "bus_number": "143", - "departure_time": "23:00:00", - "departure_seconds": now_secs + 3600, - }, - ] - ) + schedule = [ + # Already departed - must be excluded + { + "trip_id": "past_trip", + "route_id": "6656", + "bus_number": "143", + "departure_time": "00:01:00", + "departure_seconds": 60, + }, + # Future - must be included + { + "trip_id": "future_trip", + "route_id": "6656", + "bus_number": "143", + "departure_time": "23:00:00", + "departure_seconds": now_secs + 3600, + }, + ] result = get_next_departures(schedule, n=3) assert len(result) == 1 - assert result.iloc[0]["trip_id"] == "future_trip" + assert result[0]["trip_id"] == "future_trip" async def test__get_next_departures_respects_n(): @@ -192,18 +217,16 @@ async def test__get_next_departures_respects_n(): now_secs = int((now - midnight).total_seconds()) # Five future trips on the same route - n=2 should limit to 2 - schedule = pd.DataFrame( - [ - { - "trip_id": f"trip_{i}", - "route_id": "6656", - "bus_number": "143", - "departure_time": "23:00:00", - "departure_seconds": now_secs + i * 600, - } - for i in range(1, 6) - ] - ) + schedule = [ + { + "trip_id": f"trip_{i}", + "route_id": "6656", + "bus_number": "143", + "departure_time": "23:00:00", + "departure_seconds": now_secs + i * 600, + } + for i in range(1, 6) + ] assert len(get_next_departures(schedule, n=2)) == 2 assert len(get_next_departures(schedule, n=1)) == 1 @@ -214,23 +237,21 @@ async def test__get_next_departures_multiple_routes(): midnight = now.replace(hour=0, minute=0, second=0, microsecond=0) now_secs = int((now - midnight).total_seconds()) - schedule = pd.DataFrame( - [ - { - "trip_id": f"trip_{rid}_{i}", - "route_id": rid, - "bus_number": num, - "departure_time": "23:00:00", - "departure_seconds": now_secs + i * 600, - } - for rid, (_, _, num) in BUS_DATA.items() - for i in range(1, 4) - ] - ) + schedule = [ + { + "trip_id": f"trip_{rid}_{i}", + "route_id": rid, + "bus_number": num, + "departure_time": "23:00:00", + "departure_seconds": now_secs + i * 600, + } + for rid, (_, _, num) in BUS_DATA.items() + for i in range(1, 4) + ] result = get_next_departures(schedule, n=2) assert len(result) == 8 - assert set(result["route_id"]) == set(BUS_DATA.keys()) + assert {row["route_id"] for row in result} == set(BUS_DATA) # --------------------------------------------------------------------------- @@ -240,12 +261,23 @@ async def test__get_next_departures_multiple_routes(): async def test__fetch_static_schedule_returns_all_routes(): client = mock_http_client(make_gtfs_zip()) - df = await fetch_static_schedule(client) + cache = await fetch_static_schedule(client) + schedule = resolve_static_schedule(cache, datetime.now(tz=TZ_INFO).date()) - assert not df.empty + assert schedule expected_cols = {"trip_id", "route_id", "bus_number", "departure_time", "departure_seconds"} - assert expected_cols.issubset(df.columns) - assert set(df["bus_number"]) == {num for _, (_, _, num) in BUS_DATA.items()} + assert expected_cols.issubset(schedule[0]) + assert {row["bus_number"] for row in schedule} == {num for _, (_, _, num) in BUS_DATA.items()} + + +async def test__weekly_cache_resolves_multiple_weekdays_without_refetching(): + client = mock_http_client(make_gtfs_zip(active_weekdays=set(range(7)))) + cache = await fetch_static_schedule(client) + today = datetime.now(tz=TZ_INFO).date() + + assert len(resolve_static_schedule(cache, today)) == len(BUS_DATA) + assert len(resolve_static_schedule(cache, today + timedelta(days=1))) == len(BUS_DATA) + client.get.assert_awaited_once() async def test__fetch_static_schedule_excludes_wrong_direction(): @@ -275,8 +307,8 @@ async def test__fetch_static_schedule_excludes_wrong_direction(): ) client = mock_http_client(buf.getvalue()) - df = await fetch_static_schedule(client) - assert df.empty + with pytest.raises(RuntimeError, match="no departures"): + await fetch_static_schedule(client) async def test__fetch_static_schedule_raises_on_http_error(): @@ -289,10 +321,23 @@ async def test__fetch_static_schedule_raises_on_http_error(): await fetch_static_schedule(client) +async def test__fetch_static_schedule_raises_on_http_error_status(): + client = AsyncMock(spec=AsyncClient) + client.get = AsyncMock( + return_value=Response( + status_code=500, + request=Request("GET", "https://gtfs-static.translink.ca/gtfs/google_transit.zip"), + ) + ) + + with pytest.raises(RuntimeError, match="Failed to fetch static schedule"): + await fetch_static_schedule(client) + + async def test__fetch_static_schedule_raises_on_bad_zip(): client = mock_http_client(b"this is not a zip") - with pytest.raises(RuntimeError, match="Failed to read static schedule zip file"): + with pytest.raises(RuntimeError, match="Failed to parse static schedule"): await fetch_static_schedule(client) @@ -460,12 +505,11 @@ async def test__get_or_fetch_realtime_feed_returns_none_on_http_error_status(): # --------------------------------------------------------------------------- -# Tests for get_or_fetch_static_schedule +# Tests for the preprocessed static schedule cache # --------------------------------------------------------------------------- -async def test__get_or_fetch_static_schedule_cache_hit(): - """When the DB has today's row, no HTTP call should be made.""" +async def test__get_static_schedule_cache_hit(): today = datetime.now(tz=TZ_INFO).date() cached_records = [ { @@ -476,46 +520,100 @@ async def test__get_or_fetch_static_schedule_cache_hit(): "departure_seconds": 82800, } ] - cached_row = TransLinkStaticScheduleDB(id=1, date_fetched=today, schedule=cached_records) + cached_row = TransLinkStaticScheduleDB( + id=1, + date_fetched=today, + schedule=make_static_cache(cached_records, today), + ) session = mock_db_session(cached_row=cached_row) - client = AsyncMock(spec=AsyncClient) - result_date, result_df = await get_or_fetch_static_schedule(session, client) + result_date, result_rows = await get_static_schedule(session) assert result_date == today - assert not result_df.empty - assert result_df.iloc[0]["bus_number"] == "143" - client.get.assert_not_called() + assert result_rows[0]["bus_number"] == "143" + session.merge.assert_not_called() + session.commit.assert_not_called() -async def test__get_or_fetch_static_schedule_cache_miss_fetches(): - """On a cache miss the function should call the API and persist the result.""" - today = datetime.now(tz=TZ_INFO).date() +async def test__get_static_schedule_cache_miss_raises(): + session = mock_db_session(cached_row=None) + + with pytest.raises(StaticScheduleCacheUnavailableError, match=STATIC_CACHE_UNAVAILABLE_MESSAGE): + await get_static_schedule(session) + + session.merge.assert_not_called() + + +async def test__refresh_static_schedule_persists_preprocessed_cache(): session = mock_db_session(cached_row=None) client = mock_http_client(make_gtfs_zip()) - result_date, result_df = await get_or_fetch_static_schedule(session, client) + result = await refresh_static_schedule(session, client) - assert result_date == today - assert not result_df.empty + assert result["version"] == STATIC_CACHE_VERSION session.merge.assert_awaited_once() session.commit.assert_awaited_once() -async def test__get_or_fetch_static_schedule_db_write_failure_still_returns(): - """If the DB write fails, the function should still return the fetched data.""" +async def test__refresh_static_schedule_db_failure_rolls_back(): import sqlalchemy.exc - today = datetime.now(tz=TZ_INFO).date() session = mock_db_session(cached_row=None) session.merge = AsyncMock(side_effect=sqlalchemy.exc.SQLAlchemyError("disk full")) client = mock_http_client(make_gtfs_zip()) - result_date, result_df = await get_or_fetch_static_schedule(session, client) + with pytest.raises(RuntimeError, match="Failed to store static schedule"): + await refresh_static_schedule(session, client) - assert result_date == today - assert not result_df.empty session.rollback.assert_awaited_once() + session.commit.assert_not_called() + + +async def test__resolve_static_schedule_applies_calendar_exceptions(): + service_date = date(2026, 8, 13) + date_str = service_date.strftime("%Y%m%d") + regular = { + "trip_id": "regular", + "route_id": "6656", + "bus_number": "143", + "departure_time": "10:00:00", + "departure_seconds": 36000, + } + replacement = {**regular, "trip_id": "replacement", "departure_time": "11:00:00", "departure_seconds": 39600} + cache = make_static_cache([regular], service_date) + cache["services"]["SPECIAL"] = { + "start_date": date_str, + "end_date": date_str, + "weekdays": [], + } + cache["departures"]["SPECIAL"] = [replacement] + cache["exceptions"][date_str] = {"added": ["SPECIAL"], "removed": ["SVC1"]} + + assert resolve_static_schedule(cache, service_date) == [replacement] + + +async def test__resolve_static_schedule_rejects_incompatible_version(): + service_date = date(2026, 8, 13) + cache = make_static_cache([], service_date, version=STATIC_CACHE_VERSION + 1) + + with pytest.raises(StaticScheduleCacheUnavailableError, match=STATIC_CACHE_UNAVAILABLE_MESSAGE): + resolve_static_schedule(cache, service_date) + + +async def test__resolve_static_schedule_rejects_malformed_departure(): + service_date = date(2026, 8, 13) + cache = make_static_cache([{"trip_id": "missing required fields"}], service_date) + + with pytest.raises(StaticScheduleCacheUnavailableError, match=STATIC_CACHE_UNAVAILABLE_MESSAGE): + resolve_static_schedule(cache, service_date) + + +async def test__resolve_static_schedule_rejects_date_outside_coverage(): + cache_date = date(2026, 8, 13) + cache = make_static_cache([], cache_date) + + with pytest.raises(StaticScheduleCacheUnavailableError, match=STATIC_CACHE_UNAVAILABLE_MESSAGE): + resolve_static_schedule(cache, cache_date + timedelta(days=1)) async def test__get_departure_statuses_uses_timestamps_when_realtime_unavailable(): @@ -525,15 +623,18 @@ async def test__get_departure_statuses_uses_timestamps_when_realtime_unavailable cached_row = TransLinkStaticScheduleDB( id=1, date_fetched=now.date(), - schedule=[ - { - "trip_id": "trip_143", - "route_id": "6656", - "bus_number": "143", - "departure_time": "23:00:00", - "departure_seconds": departure_seconds, - } - ], + schedule=make_static_cache( + [ + { + "trip_id": "trip_143", + "route_id": "6656", + "bus_number": "143", + "departure_time": "23:00:00", + "departure_seconds": departure_seconds, + } + ], + now.date(), + ), ) session = mock_db_session() session.scalar = AsyncMock(side_effect=[cached_row, None, None]) @@ -582,21 +683,19 @@ async def test__endpoint_realtime_returns_200(client): async def test__endpoint_static_returns_schedule(client): today = datetime.now(tz=TZ_INFO).date() - mock_df = pd.DataFrame( - [ - { - "trip_id": f"trip_{num}", - "route_id": rid, - "bus_number": num, - "departure_time": "23:00:00", - "departure_seconds": 82800, - } - for rid, (_, _, num) in BUS_DATA.items() - ] - ) + mock_rows = [ + { + "trip_id": f"trip_{num}", + "route_id": rid, + "bus_number": num, + "departure_time": "23:00:00", + "departure_seconds": 82800, + } + for rid, (_, _, num) in BUS_DATA.items() + ] with patch( - "translink.urls.get_or_fetch_static_schedule", - return_value=(today, mock_df), + "translink.urls.get_static_schedule", + return_value=(today, mock_rows), ) as mock_fn: response = await client.get("/translink/static") @@ -607,6 +706,17 @@ async def test__endpoint_static_returns_schedule(client): mock_fn.assert_awaited_once() +async def test__endpoint_static_returns_503_when_cache_unavailable(client): + with patch( + "translink.urls.get_static_schedule", + side_effect=StaticScheduleCacheUnavailableError(STATIC_CACHE_UNAVAILABLE_MESSAGE), + ): + response = await client.get("/translink/static") + + assert response.status_code == status.HTTP_503_SERVICE_UNAVAILABLE + assert response.json() == {"detail": STATIC_CACHE_UNAVAILABLE_MESSAGE} + + async def test__endpoint_schedule_returns_departure_list(client): mock_results = [ TransLinkScheduleResponse( @@ -653,3 +763,14 @@ async def test__endpoint_schedule_on_time_when_no_realtime(client): assert response.status_code == status.HTTP_200_OK data = response.json() assert all(d["delay_seconds"] == 0 for d in data) + + +async def test__endpoint_schedule_returns_503_when_cache_unavailable(client): + with patch( + "translink.urls.get_departure_statuses", + side_effect=StaticScheduleCacheUnavailableError(STATIC_CACHE_UNAVAILABLE_MESSAGE), + ): + response = await client.get("/translink/schedule") + + assert response.status_code == status.HTTP_503_SERVICE_UNAVAILABLE + assert response.json() == {"detail": STATIC_CACHE_UNAVAILABLE_MESSAGE}