planet/backend/tests/test_visualization_compute_centers.py

from datetime import datetime, timezone

import pytest
from httpx import ASGITransport, AsyncClient

from app.api.v1.visualization import convert_compute_centers_to_geojson
from app.db.session import get_db
from app.main import app
from app.models.collected_data import CollectedData


def _build_record(
    *,
    record_id: int,
    source: str,
    data_type: str,
    name: str,
    country: str,
    city: str,
    latitude: float,
    longitude: float,
    metadata: dict,
):
    return CollectedData(
        id=record_id,
        source=source,
        data_type=data_type,
        source_id=f"{source}-{record_id}",
        name=name,
        extra_data={
            "country": country,
            "city": city,
            "latitude": latitude,
            "longitude": longitude,
            **metadata,
        },
        collected_at=datetime(2026, 4, 22, tzinfo=timezone.utc),
        reference_date=datetime(2026, 4, 21, tzinfo=timezone.utc),
        is_current=True,
    )


def test_convert_compute_centers_to_geojson_unifies_sources():
    top500_record = _build_record(
        record_id=1,
        source="top500",
        data_type="supercomputer",
        name="Frontier",
        country="United States",
        city="Oak Ridge",
        latitude=35.93,
        longitude=-84.31,
        metadata={
            "rank": 1,
            "manufacturer": "HPE",
            "organization": "ORNL",
            "rmax": 1102000.0,
            "cores": 8730112,
            "power": 21510.0,
        },
    )
    gpu_record = _build_record(
        record_id=2,
        source="epoch_ai_gpu",
        data_type="gpu_cluster",
        name="Colossus",
        country="United States",
        city="Memphis",
        latitude=35.15,
        longitude=-90.05,
        metadata={
            "organization": "xAI",
            "gpu_type": "H100",
            "gpu_count": 100000,
            "value": "20000",
            "unit": "TFlop/s",
        },
    )

    payload = convert_compute_centers_to_geojson([top500_record, gpu_record])

    assert payload["type"] == "FeatureCollection"
    assert len(payload["features"]) == 2

    supercomputer_feature = payload["features"][0]
    assert supercomputer_feature["properties"]["site_type"] == "supercomputer"
    assert supercomputer_feature["properties"]["capacity_unit"] == "GFlops"
    assert supercomputer_feature["properties"]["capacity_band"] == "exascale"
    assert supercomputer_feature["properties"]["operator"] == "ORNL"
    assert supercomputer_feature["properties"]["location_precision"] == "precise"
    assert supercomputer_feature["properties"]["is_estimated"] is False

    gpu_feature = payload["features"][1]
    assert gpu_feature["properties"]["site_type"] == "gpu_cluster"
    assert gpu_feature["properties"]["vendor"] == "H100"
    assert gpu_feature["properties"]["gpu_count"] == 100000
    assert gpu_feature["properties"]["capacity_band"] == "large"
    assert gpu_feature["properties"]["location_precision"] == "precise"


def test_convert_compute_centers_to_geojson_uses_coordinate_hints():
    hinted_record = _build_record(
        record_id=3,
        source="top500",
        data_type="supercomputer",
        name="Frontier",
        country="United States",
        city="",
        latitude=0.0,
        longitude=0.0,
        metadata={
            "organization": "Oak Ridge National Laboratory",
            "rmax": 1102000.0,
        },
    )

    payload = convert_compute_centers_to_geojson([hinted_record])

    assert len(payload["features"]) == 1
    coords = payload["features"][0]["geometry"]["coordinates"]
    assert coords[0] == pytest.approx(-84.3107)
    assert coords[1] == pytest.approx(35.9319)
    assert payload["features"][0]["properties"]["is_estimated"] is True
    assert payload["features"][0]["properties"]["location_precision"] == "estimated_site"


def test_convert_compute_centers_to_geojson_falls_back_to_country_centroid():
    centroid_record = _build_record(
        record_id=4,
        source="epoch_ai_gpu",
        data_type="gpu_cluster",
        name="Unknown Cluster",
        country="United States",
        city="",
        latitude=0.0,
        longitude=0.0,
        metadata={
            "organization": "Unknown Operator",
            "value": "10000",
            "unit": "TFlop/s",
        },
    )

    payload = convert_compute_centers_to_geojson([centroid_record])

    assert len(payload["features"]) == 1
    props = payload["features"][0]["properties"]
    coords = payload["features"][0]["geometry"]["coordinates"]
    assert coords[0] == pytest.approx(-98.5795)
    assert coords[1] == pytest.approx(39.8283)
    assert props["is_estimated"] is True
    assert props["location_precision"] == "estimated_country"
    assert props["geography_mode"] == "country_centroid"


@pytest.mark.asyncio
async def test_compute_centers_geojson_endpoint_returns_stats():
    records = [
        _build_record(
            record_id=1,
            source="top500",
            data_type="supercomputer",
            name="Frontier",
            country="United States",
            city="Oak Ridge",
            latitude=35.93,
            longitude=-84.31,
            metadata={"rank": 1, "rmax": 1102000.0},
        ),
        _build_record(
            record_id=2,
            source="epoch_ai_gpu",
            data_type="gpu_cluster",
            name="Colossus",
            country="United States",
            city="Memphis",
            latitude=35.15,
            longitude=-90.05,
            metadata={"value": "20000", "unit": "TFlop/s"},
        ),
    ]

    class _ScalarResult:
        def __init__(self, rows):
            self._rows = rows

        def scalars(self):
            class _Scalars:
                def __init__(self, rows):
                    self._rows = rows

                def all(self):
                    return self._rows

            return _Scalars(self._rows)

    class _FakeSession:
        async def execute(self, _query):
            return _ScalarResult(records)

    async def override_get_db():
        yield _FakeSession()

    app.dependency_overrides[get_db] = override_get_db
    transport = ASGITransport(app=app)
    try:
        async with AsyncClient(transport=transport, base_url="http://test") as client:
            response = await client.get("/api/v1/visualization/geo/compute-centers")

        assert response.status_code == 200
        data = response.json()
        assert data["count"] == 2
        assert data["stats"]["supercomputers"] == 1
        assert data["stats"]["gpu_clusters"] == 1
        assert data["features"][0]["properties"]["data_type"] == "compute_center"
    finally:
        app.dependency_overrides.clear()