1103 lines
36 KiB
Python
1103 lines
36 KiB
Python
from datetime import datetime, timezone
|
|
from unittest.mock import AsyncMock
|
|
|
|
import pytest
|
|
from httpx import ASGITransport, AsyncClient
|
|
|
|
from app.api.v1 import visualization as visualization_api
|
|
from app.api.v1.visualization import (
|
|
CollectComputeCenterLocationRequest,
|
|
convert_compute_centers_to_geojson,
|
|
)
|
|
import app.services.compute_center_locations as compute_center_locations
|
|
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
|
|
assert supercomputer_feature["properties"]["location_source"] == "source_coordinates"
|
|
assert supercomputer_feature["properties"]["location_confidence"] == 1.0
|
|
|
|
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_accepts_source_coordinate_aliases():
|
|
record = _build_record(
|
|
record_id=3,
|
|
source="epoch_ai_gpu",
|
|
data_type="gpu_cluster",
|
|
name="Alias Coordinates",
|
|
country="United States",
|
|
city="New York",
|
|
latitude=0.0,
|
|
longitude=0.0,
|
|
metadata={
|
|
"latitude": "",
|
|
"longitude": "",
|
|
"location": {
|
|
"lat": 40.7128,
|
|
"lng": -74.0060,
|
|
},
|
|
"value": "1200",
|
|
"unit": "TFlop/s",
|
|
},
|
|
)
|
|
|
|
payload = convert_compute_centers_to_geojson([record])
|
|
|
|
assert len(payload["features"]) == 1
|
|
feature = payload["features"][0]
|
|
assert feature["geometry"]["coordinates"] == [-74.006, 40.7128]
|
|
assert feature["properties"]["location_source"] == "source_coordinates"
|
|
|
|
|
|
def test_compute_center_source_coordinates_win_over_stored_location():
|
|
compute_center_locations.set_compute_center_location_cache({
|
|
"top500:top500-31": {
|
|
"source": "top500",
|
|
"source_id": "top500-31",
|
|
"name": "Stored Wrong",
|
|
"latitude": 1.0,
|
|
"longitude": 2.0,
|
|
"precision": "city",
|
|
"confidence": 0.5,
|
|
"needs_confirmation": True,
|
|
}
|
|
})
|
|
record = _build_record(
|
|
record_id=31,
|
|
source="top500",
|
|
data_type="supercomputer",
|
|
name="Source Wins",
|
|
country="United States",
|
|
city="Oak Ridge",
|
|
latitude=35.93,
|
|
longitude=-84.31,
|
|
metadata={"organization": "ORNL"},
|
|
)
|
|
|
|
payload = convert_compute_centers_to_geojson([record])
|
|
|
|
assert payload["features"][0]["geometry"]["coordinates"] == [-84.31, 35.93]
|
|
assert payload["features"][0]["properties"]["location_source"] == "source_coordinates"
|
|
compute_center_locations.set_compute_center_location_cache({})
|
|
|
|
|
|
def test_compute_center_geojson_uses_stored_location_when_source_coords_missing():
|
|
compute_center_locations.set_compute_center_location_cache({
|
|
"epoch_ai_gpu:epoch_ai_gpu-32": {
|
|
"source": "epoch_ai_gpu",
|
|
"source_id": "epoch_ai_gpu-32",
|
|
"name": "Stored Cluster",
|
|
"city": "Memphis",
|
|
"country": "United States",
|
|
"latitude": 35.1495,
|
|
"longitude": -90.049,
|
|
"precision": "city",
|
|
"confidence": 0.72,
|
|
"location_source": "manual_selection",
|
|
"source_note": "Saved by user",
|
|
"needs_confirmation": False,
|
|
"verified_at": "2026-05-08T00:00:00Z",
|
|
}
|
|
})
|
|
record = _build_record(
|
|
record_id=32,
|
|
source="epoch_ai_gpu",
|
|
data_type="gpu_cluster",
|
|
name="Stored Cluster",
|
|
country="United States",
|
|
city="",
|
|
latitude=0.0,
|
|
longitude=0.0,
|
|
metadata={"value": "1200", "unit": "TFlop/s"},
|
|
)
|
|
|
|
payload = convert_compute_centers_to_geojson([record])
|
|
|
|
assert len(payload["features"]) == 1
|
|
feature = payload["features"][0]
|
|
assert feature["geometry"]["coordinates"] == [-90.049, 35.1495]
|
|
assert feature["properties"]["location_source"] == "stored_compute_center_location"
|
|
assert feature["properties"]["needs_confirmation"] is False
|
|
compute_center_locations.set_compute_center_location_cache({})
|
|
|
|
|
|
def test_convert_compute_centers_to_geojson_does_not_use_registry_aliases():
|
|
registry_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([registry_record])
|
|
|
|
assert payload["features"] == []
|
|
assert len(payload["unresolved"]) == 1
|
|
assert payload["unresolved"][0]["name"] == "Frontier"
|
|
assert "source coords" in payload["unresolved"][0]["failure_reason"]
|
|
|
|
|
|
def test_convert_compute_centers_to_geojson_does_not_use_city_fallback():
|
|
city_record = _build_record(
|
|
record_id=4,
|
|
source="epoch_ai_gpu",
|
|
data_type="gpu_cluster",
|
|
name="Sample GPU Cluster",
|
|
country="United States",
|
|
city="San Francisco, CA",
|
|
latitude=0.0,
|
|
longitude=0.0,
|
|
metadata={
|
|
"organization": "Sample Operator",
|
|
"value": "10000",
|
|
"unit": "TFlop/s",
|
|
},
|
|
)
|
|
|
|
payload = convert_compute_centers_to_geojson([city_record])
|
|
|
|
assert payload["features"] == []
|
|
assert len(payload["unresolved"]) == 1
|
|
assert payload["unresolved"][0]["city"] == "San Francisco, CA"
|
|
|
|
|
|
def test_convert_compute_centers_to_geojson_does_not_online_geocode_on_startup(monkeypatch):
|
|
compute_center_locations._geocode_online.cache_clear()
|
|
|
|
def _explode(_query):
|
|
raise AssertionError("startup GeoJSON must not call online geocoding")
|
|
|
|
monkeypatch.setattr(compute_center_locations, "_geocode_online", _explode)
|
|
country_record = _build_record(
|
|
record_id=4,
|
|
source="epoch_ai_gpu",
|
|
data_type="gpu_cluster",
|
|
name="Unknown Cluster",
|
|
country="France",
|
|
city="",
|
|
latitude=0.0,
|
|
longitude=0.0,
|
|
metadata={
|
|
"organization": "Unknown Operator",
|
|
"value": "10000",
|
|
"unit": "TFlop/s",
|
|
},
|
|
)
|
|
|
|
payload = convert_compute_centers_to_geojson([country_record])
|
|
|
|
assert payload["features"] == []
|
|
assert len(payload["unresolved"]) == 1
|
|
assert payload["unresolved"][0]["operator"] == "Unknown Operator"
|
|
|
|
|
|
def test_convert_compute_centers_to_geojson_records_diagnostics_when_online_geocode_fails(monkeypatch):
|
|
compute_center_locations._geocode_online.cache_clear()
|
|
monkeypatch.setattr(compute_center_locations, "_geocode_online", lambda _query: None)
|
|
country_record = _build_record(
|
|
record_id=5,
|
|
source="epoch_ai_gpu",
|
|
data_type="gpu_cluster",
|
|
name="Unknown French Cluster",
|
|
country="France",
|
|
city="",
|
|
latitude=0.0,
|
|
longitude=0.0,
|
|
metadata={
|
|
"organization": "Unknown Operator",
|
|
"value": "10000",
|
|
"unit": "TFlop/s",
|
|
},
|
|
)
|
|
|
|
payload = convert_compute_centers_to_geojson([country_record])
|
|
|
|
assert payload["features"] == []
|
|
assert len(payload["unresolved"]) == 1
|
|
diagnostic = payload["unresolved"][0]
|
|
assert diagnostic["record_id"] == 5
|
|
assert diagnostic["source_id"] == "epoch_ai_gpu-5"
|
|
assert diagnostic["country"] == "France"
|
|
assert diagnostic["operator"] == "Unknown Operator"
|
|
assert diagnostic["failure_reason"]
|
|
assert diagnostic["attempted_queries"] == []
|
|
|
|
|
|
def test_convert_compute_centers_to_geojson_records_diagnostics_when_no_country(monkeypatch):
|
|
compute_center_locations._geocode_online.cache_clear()
|
|
monkeypatch.setattr(compute_center_locations, "_geocode_online", lambda _query: None)
|
|
unknown_record = _build_record(
|
|
record_id=6,
|
|
source="epoch_ai_gpu",
|
|
data_type="gpu_cluster",
|
|
name="Unknown Offshore Cluster",
|
|
country="",
|
|
city="",
|
|
latitude=0.0,
|
|
longitude=0.0,
|
|
metadata={
|
|
"organization": "Unknown Operator",
|
|
"value": "10000",
|
|
"unit": "TFlop/s",
|
|
},
|
|
)
|
|
|
|
payload = convert_compute_centers_to_geojson([unknown_record])
|
|
|
|
assert payload["features"] == []
|
|
assert len(payload["unresolved"]) == 1
|
|
assert payload["unresolved"][0]["failure_reason"]
|
|
|
|
|
|
def test_convert_compute_centers_to_geojson_never_emits_zero_coordinates(monkeypatch):
|
|
compute_center_locations._geocode_online.cache_clear()
|
|
|
|
def _zero_geocode(query):
|
|
return {
|
|
"lat": "0",
|
|
"lon": "0",
|
|
"display_name": "Null Island",
|
|
"address": {"city": "", "country": ""},
|
|
}
|
|
|
|
monkeypatch.setattr(compute_center_locations, "_geocode_online", _zero_geocode)
|
|
record = _build_record(
|
|
record_id=7,
|
|
source="epoch_ai_gpu",
|
|
data_type="gpu_cluster",
|
|
name="Null Island Cluster",
|
|
country="",
|
|
city="",
|
|
latitude=0.0,
|
|
longitude=0.0,
|
|
metadata={"organization": "Null Inc"},
|
|
)
|
|
|
|
payload = convert_compute_centers_to_geojson([record])
|
|
|
|
for feature in payload["features"]:
|
|
coords = feature["geometry"]["coordinates"]
|
|
assert coords[0] not in (0, 0.0)
|
|
assert coords[1] not in (0, 0.0)
|
|
|
|
|
|
def test_convert_compute_centers_to_geojson_rejects_country_or_unknown_precision(monkeypatch):
|
|
compute_center_locations._geocode_online.cache_clear()
|
|
monkeypatch.setattr(compute_center_locations, "_geocode_online", lambda _query: None)
|
|
record = _build_record(
|
|
record_id=8,
|
|
source="top500",
|
|
data_type="supercomputer",
|
|
name="Phantom System",
|
|
country="Liechtenstein",
|
|
city="",
|
|
latitude=0.0,
|
|
longitude=0.0,
|
|
metadata={"organization": "Phantom Operator", "rmax": 100.0},
|
|
)
|
|
|
|
payload = convert_compute_centers_to_geojson([record])
|
|
|
|
for feature in payload["features"]:
|
|
assert feature["properties"]["location_precision"] in {"precise", "site", "city"}
|
|
assert payload["features"] == []
|
|
assert payload["unresolved"], "phantom record must surface as diagnostic"
|
|
|
|
|
|
def test_resolve_full_returns_diagnostic_for_unresolved(monkeypatch):
|
|
compute_center_locations._geocode_online.cache_clear()
|
|
monkeypatch.setattr(compute_center_locations, "_geocode_online", lambda _query: None)
|
|
record = _build_record(
|
|
record_id=11,
|
|
source="epoch_ai_gpu",
|
|
data_type="gpu_cluster",
|
|
name="Phantom Cluster",
|
|
country="Bhutan",
|
|
city="",
|
|
latitude=0.0,
|
|
longitude=0.0,
|
|
metadata={"organization": "Mystery Operator"},
|
|
)
|
|
result = compute_center_locations.resolve_compute_center_location_full(record, record.extra_data)
|
|
assert result.location is None
|
|
assert result.diagnostic is not None
|
|
assert result.diagnostic.failure_reason
|
|
assert result.diagnostic.country == "Bhutan"
|
|
|
|
|
|
def test_collect_location_candidates_ignores_registry_and_uses_online(monkeypatch):
|
|
compute_center_locations._geocode_online.cache_clear()
|
|
|
|
def _fake_ror(query):
|
|
assert query == "Oak Ridge National Laboratory"
|
|
return {
|
|
"id": "https://ror.org/01qz5mb56",
|
|
"names": [
|
|
{"types": ["ror_display"], "value": "Oak Ridge National Laboratory"}
|
|
],
|
|
"locations": [
|
|
{
|
|
"geonames_id": 4646571,
|
|
"geonames_details": {
|
|
"name": "Oak Ridge",
|
|
"country_subdivision_name": "Tennessee",
|
|
"country_name": "United States",
|
|
"lat": 36.01036,
|
|
"lng": -84.26964,
|
|
},
|
|
}
|
|
],
|
|
}
|
|
|
|
monkeypatch.setattr(compute_center_locations, "_lookup_ror_organization", _fake_ror)
|
|
monkeypatch.setattr(compute_center_locations, "_geocode_online", lambda _query: None)
|
|
candidates, attempted = compute_center_locations.collect_location_candidates(
|
|
name="Frontier",
|
|
operator="Oak Ridge National Laboratory",
|
|
country="United States",
|
|
)
|
|
assert candidates, "online source-traced query must produce a candidate"
|
|
best = candidates[0]
|
|
assert best.source == "ror_organization_registry"
|
|
assert best.precision == "city"
|
|
assert best.needs_confirmation is True
|
|
assert attempted[0] == "ror:Oak Ridge National Laboratory"
|
|
|
|
|
|
def test_collect_location_candidates_returns_online_when_registry_misses(monkeypatch):
|
|
compute_center_locations._geocode_online.cache_clear()
|
|
|
|
def _fake_geocode(query):
|
|
if "Lyon" not in query and "Mystery Operator" not in query and "Lyon, France" not in query:
|
|
return None
|
|
return {
|
|
"lat": "45.7640",
|
|
"lon": "4.8357",
|
|
"display_name": "Lyon, Auvergne-Rhône-Alpes, France",
|
|
"address": {"city": "Lyon", "state": "Auvergne-Rhône-Alpes", "country": "France"},
|
|
}
|
|
|
|
monkeypatch.setattr(compute_center_locations, "_geocode_online", _fake_geocode)
|
|
monkeypatch.setattr(compute_center_locations, "_lookup_ror_organization", lambda _query: None)
|
|
candidates, attempted = compute_center_locations.collect_location_candidates(
|
|
name="Mystery System",
|
|
operator="Mystery Operator",
|
|
city="Lyon",
|
|
country="France",
|
|
)
|
|
assert candidates, "online geocoding must produce a candidate"
|
|
online_candidates = [c for c in candidates if c.source == "nominatim_online_geocode"]
|
|
assert online_candidates, "must include at least one online candidate"
|
|
online = online_candidates[0]
|
|
assert online.precision == "city"
|
|
assert online.needs_confirmation is True
|
|
assert online.suggested_registry_entry is not None
|
|
assert attempted, "must record attempted query strings"
|
|
|
|
|
|
def test_collect_location_candidates_failure_returns_attempted_queries(monkeypatch):
|
|
compute_center_locations._geocode_online.cache_clear()
|
|
monkeypatch.setattr(compute_center_locations, "_lookup_ror_organization", lambda _query: None)
|
|
monkeypatch.setattr(compute_center_locations, "_geocode_online", lambda _query: None)
|
|
candidates, attempted = compute_center_locations.collect_location_candidates(
|
|
name="Mystery Offshore Cluster",
|
|
operator="Mystery Operator",
|
|
country="Bhutan",
|
|
)
|
|
assert candidates == []
|
|
assert attempted, "even on failure we record attempted queries for diagnostics"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_collect_compute_center_location_skips_llm_when_candidates_exist(monkeypatch):
|
|
candidate = compute_center_locations.LocationCandidate(
|
|
latitude=45.764,
|
|
longitude=4.8357,
|
|
display_name="Lyon",
|
|
precision="city",
|
|
confidence=0.62,
|
|
query="Lyon, France",
|
|
source="nominatim_online_geocode",
|
|
source_note="fixture",
|
|
matched_fields=("city", "country"),
|
|
needs_confirmation=True,
|
|
city="Lyon",
|
|
country="France",
|
|
)
|
|
monkeypatch.setattr(visualization_api, "_load_compute_center_record", AsyncMock(return_value=None))
|
|
monkeypatch.setattr(
|
|
visualization_api,
|
|
"collect_location_candidates",
|
|
lambda **_kwargs: ([candidate], ["Lyon, France"]),
|
|
)
|
|
|
|
async def _explode(**_kwargs):
|
|
raise AssertionError("LLM fallback should not run when a normal candidate exists")
|
|
|
|
monkeypatch.setattr(visualization_api, "collect_llm_location_fallback_candidate", _explode)
|
|
|
|
response = await visualization_api.collect_compute_center_location(
|
|
"epoch_ai_gpu-test",
|
|
CollectComputeCenterLocationRequest(
|
|
name="Mystery Cluster",
|
|
source="epoch_ai_gpu",
|
|
city="Lyon",
|
|
country="France",
|
|
),
|
|
db=AsyncMock(),
|
|
)
|
|
|
|
assert response["success"] is True
|
|
assert response["best_candidate"]["source"] == "nominatim_online_geocode"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_collect_compute_center_location_uses_llm_when_candidates_empty(monkeypatch):
|
|
llm_candidate = compute_center_locations.LocationCandidate(
|
|
latitude=45.764,
|
|
longitude=4.8357,
|
|
display_name="Lyon, France",
|
|
precision="city",
|
|
confidence=0.74,
|
|
query="llm_factcheck:compute_center:Mystery Cluster",
|
|
source="llm_location_factcheck",
|
|
source_note="LLM location factcheck fallback",
|
|
matched_fields=("name",),
|
|
needs_confirmation=True,
|
|
city="Lyon",
|
|
country="France",
|
|
)
|
|
monkeypatch.setattr(visualization_api, "_load_compute_center_record", AsyncMock(return_value=None))
|
|
monkeypatch.setattr(
|
|
visualization_api,
|
|
"collect_location_candidates",
|
|
lambda **_kwargs: ([], ["Mystery Cluster, France"]),
|
|
)
|
|
|
|
from app.services.location.llm_fallback import LocationLLMFallbackResult
|
|
|
|
async def _fallback(**_kwargs):
|
|
return LocationLLMFallbackResult(
|
|
candidates=[llm_candidate],
|
|
attempted_queries=["llm_factcheck:compute_center:Mystery Cluster"],
|
|
)
|
|
|
|
monkeypatch.setattr(visualization_api, "get_ai_provider_client", AsyncMock(return_value=object()))
|
|
monkeypatch.setattr(visualization_api, "collect_llm_location_fallback_candidate", _fallback)
|
|
|
|
response = await visualization_api.collect_compute_center_location(
|
|
"epoch_ai_gpu-test",
|
|
CollectComputeCenterLocationRequest(
|
|
name="Mystery Cluster",
|
|
source="epoch_ai_gpu",
|
|
country="France",
|
|
),
|
|
db=AsyncMock(),
|
|
)
|
|
|
|
assert response["success"] is True
|
|
assert response["best_candidate"]["source"] == "llm_location_factcheck"
|
|
assert response["best_candidate"]["needs_confirmation"] is True
|
|
assert response["attempted_queries"] == [
|
|
"Mystery Cluster, France",
|
|
"llm_factcheck:compute_center:Mystery Cluster",
|
|
]
|
|
|
|
|
|
@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 all(self):
|
|
return list(self._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):
|
|
query_text = str(query).lower()
|
|
if "bgp_incidents" in query_text:
|
|
return _ScalarResult(scalar_value=2)
|
|
if "bgp_anomalies" in query_text:
|
|
return _ScalarResult(scalar_value=3)
|
|
if "ais_raw_observations" in query_text or "vessel_position" in query_text:
|
|
return _ScalarResult(rows=[])
|
|
return _ScalarResult(rows=records)
|
|
|
|
async def get(self, *_args, **_kwargs):
|
|
return None
|
|
|
|
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()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_collect_location_endpoint_returns_candidates_for_known_record(monkeypatch):
|
|
def _fake_ror(query):
|
|
assert query == "Oak Ridge National Laboratory"
|
|
return {
|
|
"id": "https://ror.org/01qz5mb56",
|
|
"names": [
|
|
{"types": ["ror_display"], "value": "Oak Ridge National Laboratory"}
|
|
],
|
|
"locations": [
|
|
{
|
|
"geonames_id": 4646571,
|
|
"geonames_details": {
|
|
"name": "Oak Ridge",
|
|
"country_subdivision_name": "Tennessee",
|
|
"country_name": "United States",
|
|
"lat": 36.01036,
|
|
"lng": -84.26964,
|
|
},
|
|
}
|
|
],
|
|
}
|
|
|
|
monkeypatch.setattr(compute_center_locations, "_lookup_ror_organization", _fake_ror)
|
|
monkeypatch.setattr(compute_center_locations, "_geocode_online", lambda _query: None)
|
|
|
|
target_record = _build_record(
|
|
record_id=42,
|
|
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},
|
|
)
|
|
|
|
class _ScalarResult:
|
|
def __init__(self, rows):
|
|
self._rows = rows
|
|
|
|
def scalars(self):
|
|
class _Scalars:
|
|
def __init__(self, rows):
|
|
self._rows = rows
|
|
|
|
def first(self):
|
|
return self._rows[0] if self._rows else None
|
|
|
|
def all(self):
|
|
return self._rows
|
|
|
|
return _Scalars(self._rows)
|
|
|
|
class _FakeSession:
|
|
async def execute(self, _query):
|
|
return _ScalarResult([target_record])
|
|
|
|
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.post(
|
|
"/api/v1/visualization/compute-centers/top500-42/collect-location",
|
|
json={
|
|
"name": "Frontier",
|
|
"operator": "Oak Ridge National Laboratory",
|
|
"country": "United States",
|
|
},
|
|
)
|
|
assert response.status_code == 200
|
|
body = response.json()
|
|
assert body["success"] is True
|
|
assert body["candidates"], "must include candidates"
|
|
best = body["best_candidate"]
|
|
assert best["precision"] in {"precise", "site", "city"}
|
|
assert best["source"] == "ror_organization_registry"
|
|
assert best["needs_confirmation"] is True
|
|
assert best["matched_fields"], "matched_fields must be populated"
|
|
finally:
|
|
app.dependency_overrides.clear()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_collect_location_endpoint_returns_failure_reason(monkeypatch):
|
|
monkeypatch.setattr(compute_center_locations, "_lookup_ror_organization", lambda _query: None)
|
|
monkeypatch.setattr(compute_center_locations, "_geocode_online", lambda _query: None)
|
|
|
|
class _ScalarResult:
|
|
def __init__(self, rows):
|
|
self._rows = rows
|
|
|
|
def scalars(self):
|
|
class _Scalars:
|
|
def __init__(self, rows):
|
|
self._rows = rows
|
|
|
|
def first(self):
|
|
return self._rows[0] if self._rows else None
|
|
|
|
def all(self):
|
|
return self._rows
|
|
|
|
return _Scalars(self._rows)
|
|
|
|
class _FakeSession:
|
|
async def execute(self, _query):
|
|
return _ScalarResult([])
|
|
|
|
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.post(
|
|
"/api/v1/visualization/compute-centers/epoch-mystery-99/collect-location",
|
|
json={
|
|
"name": "Mystery Cluster",
|
|
"operator": "Mystery Operator",
|
|
"country": "Bhutan",
|
|
},
|
|
)
|
|
assert response.status_code == 200
|
|
body = response.json()
|
|
assert body["success"] is False
|
|
assert body["failure_reason"]
|
|
assert body["candidates"] == []
|
|
assert body["attempted_queries"], "must include attempted queries"
|
|
finally:
|
|
app.dependency_overrides.clear()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_save_location_endpoint_upserts_and_geojson_can_render():
|
|
target_record = _build_record(
|
|
record_id=52,
|
|
source="epoch_ai_gpu",
|
|
data_type="gpu_cluster",
|
|
name="Saved Cluster",
|
|
country="United States",
|
|
city="",
|
|
latitude=0.0,
|
|
longitude=0.0,
|
|
metadata={"value": "1200", "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 first(self):
|
|
return self._rows[0] if self._rows else None
|
|
|
|
def all(self):
|
|
return self._rows
|
|
|
|
return _Scalars(self._rows)
|
|
|
|
class _FakeSession:
|
|
def __init__(self):
|
|
self.saved = []
|
|
|
|
async def execute(self, _query):
|
|
if self.saved:
|
|
return _ScalarResult(self.saved)
|
|
return _ScalarResult([target_record])
|
|
|
|
async def scalar(self, _query):
|
|
return None
|
|
|
|
def add(self, record):
|
|
self.saved.append(record)
|
|
|
|
async def commit(self):
|
|
return None
|
|
|
|
async def refresh(self, _record):
|
|
return None
|
|
|
|
fake_session = _FakeSession()
|
|
|
|
async def override_get_db():
|
|
yield fake_session
|
|
|
|
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.post(
|
|
"/api/v1/visualization/compute-centers/epoch_ai_gpu-52/location",
|
|
json={
|
|
"source": "epoch_ai_gpu",
|
|
"name": "Saved Cluster",
|
|
"latitude": 35.1495,
|
|
"longitude": -90.049,
|
|
"precision": "city",
|
|
"confidence": 0.72,
|
|
"location_source": "ror_organization_registry",
|
|
"source_note": "Selected by user",
|
|
"raw_payload": {"source": "ror_organization_registry"},
|
|
},
|
|
)
|
|
assert response.status_code == 200
|
|
body = response.json()
|
|
assert body["success"] is True
|
|
assert fake_session.saved
|
|
|
|
payload = convert_compute_centers_to_geojson([target_record])
|
|
assert len(payload["features"]) == 1
|
|
feature = payload["features"][0]
|
|
assert feature["geometry"]["coordinates"] == [-90.049, 35.1495]
|
|
assert feature["properties"]["location_source"] == "stored_compute_center_location"
|
|
finally:
|
|
app.dependency_overrides.clear()
|
|
compute_center_locations.set_compute_center_location_cache({})
|
|
|
|
|
|
def test_resolution_chain_orders_source_coords_first(monkeypatch):
|
|
def _explode(_query):
|
|
raise AssertionError("source coords must short-circuit before online geocoding")
|
|
|
|
monkeypatch.setattr(compute_center_locations, "_geocode_online", _explode)
|
|
record = _build_record(
|
|
record_id=20,
|
|
source="top500",
|
|
data_type="supercomputer",
|
|
name="Frontier",
|
|
country="United States",
|
|
city="Oak Ridge",
|
|
latitude=35.93,
|
|
longitude=-84.31,
|
|
metadata={"organization": "ORNL"},
|
|
)
|
|
result = compute_center_locations.resolve_compute_center_location_full(record, record.extra_data)
|
|
assert result.is_resolved
|
|
assert result.location.location_precision == "precise"
|
|
assert result.location.location_source == "source_coordinates"
|
|
|
|
|
|
def test_no_country_centroid_or_major_compute_city_fallback(monkeypatch):
|
|
monkeypatch.setattr(compute_center_locations, "_geocode_online", lambda _query: None)
|
|
record = _build_record(
|
|
record_id=21,
|
|
source="top500",
|
|
data_type="supercomputer",
|
|
name="Phantom System",
|
|
country="France",
|
|
city="",
|
|
latitude=0.0,
|
|
longitude=0.0,
|
|
metadata={"organization": "Phantom Operator"},
|
|
)
|
|
result = compute_center_locations.resolve_compute_center_location_full(record, record.extra_data)
|
|
assert result.location is None, "must NOT fall back to country centroid or hashed major city"
|
|
assert result.diagnostic is not None
|
|
assert result.diagnostic.failure_reason
|
|
|
|
|
|
def test_repository_has_no_forbidden_precision_tokens():
|
|
"""Static guard: forbidden fallback strategies must not regress into the codebase.
|
|
|
|
Each forbidden token may appear at most once per target file, and only inside
|
|
the FORBIDDEN_PRECISIONS guard list (so we still reject them at runtime).
|
|
"""
|
|
from pathlib import Path
|
|
|
|
backend_root = Path(__file__).resolve().parents[1]
|
|
forbidden_tokens = (
|
|
"country_centroid",
|
|
"country_major_compute_city",
|
|
"estimated_country",
|
|
)
|
|
targets = [
|
|
backend_root / "app" / "services" / "compute_center_locations.py",
|
|
backend_root / "app" / "api" / "v1" / "visualization.py",
|
|
]
|
|
for target in targets:
|
|
text = target.read_text(encoding="utf-8")
|
|
for token in forbidden_tokens:
|
|
occurrences = text.count(token)
|
|
assert occurrences <= 1, (
|
|
f"{token} appears {occurrences} times in {target}; "
|
|
"should only appear in FORBIDDEN_PRECISIONS guard list."
|
|
)
|
|
if occurrences == 1:
|
|
assert "FORBIDDEN_PRECISIONS" in text, (
|
|
f"{token} appears in {target} outside the FORBIDDEN_PRECISIONS guard"
|
|
)
|