348 lines
11 KiB
Python
348 lines
11 KiB
Python
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()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_visualization_geo_summary_returns_counts(monkeypatch):
|
|
records = [
|
|
_build_record(
|
|
record_id=1,
|
|
source="arcgis_cables",
|
|
data_type="submarine_cable",
|
|
name="Test Cable",
|
|
country="",
|
|
city="",
|
|
latitude=0,
|
|
longitude=0,
|
|
metadata={
|
|
"route_coordinates": [[[0, 0], [1, 1]]],
|
|
"status": "active",
|
|
},
|
|
),
|
|
_build_record(
|
|
record_id=2,
|
|
source="arcgis_landing_points",
|
|
data_type="landing_point",
|
|
name="Test Landing",
|
|
country="United States",
|
|
city="New York",
|
|
latitude=40.7,
|
|
longitude=-74.0,
|
|
metadata={"city_id": 10},
|
|
),
|
|
_build_record(
|
|
record_id=3,
|
|
source="celestrak_tle",
|
|
data_type="satellite_tle",
|
|
name="TESTSAT",
|
|
country="",
|
|
city="",
|
|
latitude=0,
|
|
longitude=0,
|
|
metadata={
|
|
"norad_cat_id": 12345,
|
|
"tle_line1": "1 12345U 98067A 24001.00000000 .00000000 00000-0 00000-0 0 9991",
|
|
"tle_line2": "2 12345 51.6000 100.0000 0001000 10.0000 20.0000 15.50000000 01",
|
|
},
|
|
),
|
|
_build_record(
|
|
record_id=4,
|
|
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=5,
|
|
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=None, scalar_value=None):
|
|
self._rows = rows or []
|
|
self._scalar_value = scalar_value
|
|
|
|
def scalar(self):
|
|
return self._scalar_value
|
|
|
|
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):
|
|
query_text = str(query)
|
|
if "bgp_incidents" in query_text:
|
|
return _ScalarResult(scalar_value=2)
|
|
if "bgp_anomalies" in query_text:
|
|
return _ScalarResult(scalar_value=3)
|
|
return _ScalarResult(rows=records)
|
|
|
|
async def override_get_db():
|
|
yield _FakeSession()
|
|
|
|
async def _fake_build_bgp_collector_coverage(*_args, **_kwargs):
|
|
return [
|
|
{"collector": "rrc00"},
|
|
{"collector": "rrc01"},
|
|
]
|
|
|
|
monkeypatch.setattr(
|
|
"app.api.v1.visualization.build_bgp_collector_coverage",
|
|
_fake_build_bgp_collector_coverage,
|
|
)
|
|
|
|
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/summary")
|
|
|
|
assert response.status_code == 200
|
|
stats = response.json()["stats"]
|
|
assert stats["cable_count"] == 1
|
|
assert stats["landing_point_count"] == 1
|
|
assert stats["satellite_count"] == 1
|
|
assert stats["compute_center_count"] == 2
|
|
assert stats["supercomputer_count"] == 1
|
|
assert stats["gpu_cluster_count"] == 1
|
|
assert stats["bgp_event_count"] == 2
|
|
assert stats["bgp_incident_count"] == 2
|
|
assert stats["bgp_anomaly_count"] == 3
|
|
assert stats["bgp_collector_count"] == 2
|
|
finally:
|
|
app.dependency_overrides.clear()
|