planet/backend/tests/test_visualization_satellite_streaming.py

from collections.abc import AsyncIterator, Iterator, Sequence
from datetime import UTC, datetime, timedelta

import pytest
from sqlalchemy import Select, create_engine
from sqlalchemy.engine import Result, ScalarResult
from sqlalchemy.orm import Session

from app.api.v1 import visualization
from app.models.collected_data import CollectedData


class SatelliteRows:
    def __init__(self, result: ScalarResult[CollectedData]) -> None:
        self.result = result
        self.closed = False

    async def partitions(self, size: int) -> AsyncIterator[Sequence[CollectedData]]:
        for batch in self.result.partitions(size):
            yield batch

    async def close(self) -> None:
        self.result.close()
        self.closed = True


class SatelliteSession:
    def __init__(self, session: Session) -> None:
        self.session = session
        self.streams: list[SatelliteRows] = []

    async def stream_scalars(
        self, statement: Select, *, execution_options: dict[str, int]
    ) -> SatelliteRows:
        rows = SatelliteRows(self.session.scalars(statement, execution_options=execution_options))
        self.streams.append(rows)
        return rows

    async def execute(self, statement: Select) -> Result:
        return self.session.execute(statement)


@pytest.fixture
def satellite_db() -> Iterator[SatelliteSession]:
    engine = create_engine("sqlite://")
    CollectedData.__table__.create(engine)
    with Session(engine) as session:
        yield SatelliteSession(session)
    engine.dispose()


def satellite(
    record_id: int, *, current: bool = True, task_id: int = 1,
    name: str | None = None, source: str = "celestrak_tle",
) -> CollectedData:
    return CollectedData(
        id=record_id,
        source=source,
        source_id=str(record_id),
        data_type="satellite_tle",
        name=name or f"SAT-{record_id}",
        is_current=current,
        task_id=task_id,
        collected_at=datetime(2026, 1, 1, tzinfo=UTC) + timedelta(days=task_id),
        extra_data={
            "norad_cat_id": record_id + 40000,
            "tle_line1": f"line-1-{record_id}",
            "tle_line2": f"line-2-{record_id}",
        },
    )


@pytest.mark.asyncio
@pytest.mark.parametrize("limit", [None, 1, 501])
async def test_streamed_satellites_preserve_all_rows_order_tle_and_limit(
    satellite_db: SatelliteSession, limit: int | None,
) -> None:
    satellite_db.session.add_all([satellite(record_id) for record_id in range(1, 1004)])
    satellite_db.session.add_all([
        satellite(2000, name="Unknown"),
        satellite(3000, current=False),
        satellite(4000, source="other"),
    ])
    satellite_db.session.commit()

    payload = await visualization._build_satellites_geojson(limit=limit, db=satellite_db)

    expected_ids = list(range(1003, 0, -1))[:limit]
    assert payload["count"] == len(expected_ids)
    assert [feature["properties"]["id"] for feature in payload["features"]] == expected_ids
    assert [feature["properties"]["tle_line1"] for feature in payload["features"]] == [
        f"line-1-{record_id}" for record_id in expected_ids
    ]
    assert len(satellite_db.streams) == 1
    assert satellite_db.streams[0].closed


@pytest.mark.asyncio
async def test_streamed_satellites_fall_back_to_latest_non_unknown_task(
    satellite_db: SatelliteSession,
) -> None:
    satellite_db.session.add_all([
        satellite(1, current=False, task_id=1),
        satellite(2, current=False, task_id=2),
        satellite(3, current=False, task_id=2),
        satellite(4, current=False, task_id=3, name="Unknown"),
    ])
    satellite_db.session.commit()

    payload = await visualization._build_satellites_geojson(limit=1, db=satellite_db)

    assert payload["count"] == 1
    assert payload["features"][0]["properties"]["id"] == 3
    assert all(stream.closed for stream in satellite_db.streams)


@pytest.mark.asyncio
async def test_unrenderable_current_rows_do_not_restore_historical_satellites(
    satellite_db: SatelliteSession,
) -> None:
    current = satellite(2)
    current.extra_data = {}
    satellite_db.session.add_all([satellite(1, current=False), current])
    satellite_db.session.commit()

    payload = await visualization._build_satellites_geojson(limit=None, db=satellite_db)

    assert payload == {"type": "FeatureCollection", "features": [], "count": 0}
    assert len(satellite_db.streams) == 1


@pytest.mark.asyncio
async def test_empty_satellite_catalog_returns_empty_collection(
    satellite_db: SatelliteSession,
) -> None:
    payload = await visualization._build_satellites_geojson(limit=None, db=satellite_db)

    assert payload == {"type": "FeatureCollection", "features": [], "count": 0}
    assert satellite_db.streams[0].closed


@pytest.mark.asyncio
async def test_satellite_cursor_closes_if_conversion_fails(
    satellite_db: SatelliteSession, monkeypatch: pytest.MonkeyPatch,
) -> None:
    satellite_db.session.add(satellite(1))
    satellite_db.session.commit()

    def fail_conversion(_records: Sequence[CollectedData]) -> None:
        raise ValueError("invalid orbital data")

    monkeypatch.setattr(visualization, "convert_satellite_to_geojson", fail_conversion)
    with pytest.raises(ValueError, match="invalid orbital data"):
        await visualization._build_satellites_geojson(limit=None, db=satellite_db)

    assert satellite_db.streams[0].closed