forked from linkong/planet
269 lines
11 KiB
Python
269 lines
11 KiB
Python
import csv
|
|
import io
|
|
from datetime import UTC, datetime
|
|
from unittest.mock import AsyncMock, MagicMock
|
|
|
|
import pytest
|
|
from fastapi import HTTPException
|
|
from sqlalchemy.dialects import postgresql
|
|
|
|
from app.api.v1.compute_intelligence import sample_query, selected_countries
|
|
from app.services.collectors.epoch_ai import EpochAIGPUCollector
|
|
from app.services.collectors.top500 import TOP500Collector
|
|
from app.services.compute_history import compute_history
|
|
|
|
|
|
def epoch_csv(**overrides):
|
|
row = {
|
|
"Name": "Cluster Alpha",
|
|
"Country": "United States of America",
|
|
"Status": "Existing",
|
|
"First Operational Date": "2023-05-02",
|
|
"16-bit OP/s (log)": "18",
|
|
"Location": "3231 Riverport Rd, Memphis, TN 38109",
|
|
"latitude": "35",
|
|
"longitude": "-90",
|
|
**overrides,
|
|
}
|
|
output = io.StringIO()
|
|
writer = csv.DictWriter(output, fieldnames=list(row))
|
|
writer.writeheader()
|
|
writer.writerow(row)
|
|
return output.getvalue()
|
|
|
|
|
|
def test_epoch_csv_preserves_precision_date_and_missing_values():
|
|
collector = EpochAIGPUCollector()
|
|
row = collector.parse_response(epoch_csv())[0]
|
|
assert row["value"] == 1e18
|
|
assert row["unit"] == "FLOP/s"
|
|
assert row["metadata"]["precision"] == "FP16"
|
|
assert row["metadata"]["location_precision"] == "precise"
|
|
assert row["metadata"]["raw_country"] == "United States of America"
|
|
assert row["reference_date"] == "2023-05-02"
|
|
missing = collector.parse_response(
|
|
epoch_csv(**{"16-bit OP/s (log)": "", "First Operational Date": ""})
|
|
)[0]
|
|
assert missing["value"] is None
|
|
assert missing["reference_date"] is None
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("location", "country", "expected"),
|
|
[
|
|
("United States of America", "United States of America", "country"),
|
|
("Japan", "Japan", "country"),
|
|
("Korea (Republic of)", "South Korea", "country"),
|
|
("Texas", "United States of America", "region"),
|
|
("Iowa, USA", "United States of America", "region"),
|
|
("Seoul", "South Korea", "city"),
|
|
("Paris, France", "France", "city"),
|
|
("Skolkovo Innovation Center\nMoscow Oblast\nRussia", "Russia", "site"),
|
|
("Cluster 2, Texas", "United States of America", "site"),
|
|
("Campus 7, Tokyo", "Japan", "site"),
|
|
("Building 240, Argonne", "United States of America", "site"),
|
|
("7", "United States of America", "unknown"),
|
|
("Jülich Supercomputing Centre\n52428 Jülich, Germany", "Germany", "precise"),
|
|
("35.380924,139.625008", "Japan", "precise"),
|
|
('33°24\'28.7"N 112°21\'54.3"W', "United States of America", "precise"),
|
|
("latitude: 35.0 longitude: 181.0", "United States of America", "unknown"),
|
|
('99°99\'00"N 200°00\'00"W', "United States of America", "unknown"),
|
|
("3231 Riverport Rd, Memphis, TN 38109", "United States of America", "precise"),
|
|
],
|
|
)
|
|
def test_epoch_location_precision_follows_location_field(location, country, expected):
|
|
row = EpochAIGPUCollector().parse_response(
|
|
epoch_csv(Location=location, Country=country)
|
|
)[0]
|
|
|
|
assert row["metadata"]["location_precision"] == expected
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("location", "expected"),
|
|
[
|
|
({"lat": 35.0, "lng": 139.0}, "precise"),
|
|
({"latitude": 35.0, "longitude": 181.0}, "unknown"),
|
|
({"lat": 35.0}, "unknown"),
|
|
({"lng": 139.0}, "unknown"),
|
|
({"lat": "invalid", "lng": 139.0}, "unknown"),
|
|
],
|
|
)
|
|
def test_epoch_coordinate_object_requires_complete_valid_pair(location, expected):
|
|
from app.services.epoch_location import assess_epoch_location
|
|
|
|
assert assess_epoch_location(location).precision == expected
|
|
|
|
|
|
def test_epoch_anonymous_rows_are_not_geocoded():
|
|
row = EpochAIGPUCollector().parse_response(epoch_csv(Name="Anonymized 1", Country="China"))[0]
|
|
assert row["latitude"] is None and row["longitude"] is None
|
|
assert row["metadata"]["anonymous"] and row["metadata"]["rounded"]
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"payload",
|
|
["<html>blocked</html>", "Name,Country\nA,US", epoch_csv(**{"16-bit OP/s (log)": "nan"})],
|
|
)
|
|
def test_epoch_invalid_payload_fails_closed(payload):
|
|
with pytest.raises(ValueError):
|
|
EpochAIGPUCollector().parse_response(payload)
|
|
|
|
|
|
def test_epoch_duplicate_identity_rejected():
|
|
payload = epoch_csv()
|
|
with pytest.raises(ValueError, match="identity collision"):
|
|
EpochAIGPUCollector().parse_response(payload + payload.splitlines()[1] + "\n")
|
|
|
|
|
|
def test_anonymous_source_rows_preserve_multiplicity_without_claiming_identity():
|
|
payload = epoch_csv(Name="Anonymized Chinese System", Country="China")
|
|
rows = EpochAIGPUCollector().parse_response(payload + payload.splitlines()[1] + "\n")
|
|
assert len(rows) == 2
|
|
assert rows[0]["source_id"] != rows[1]["source_id"]
|
|
assert all(row["metadata"]["identity_scope"] == "snapshot_record" for row in rows)
|
|
|
|
|
|
def test_top500_units_and_stable_identity():
|
|
collector = TOP500Collector()
|
|
html = '<table id="top500"><tr><th>Rank</th><th>System</th><th>Cores</th><th>Rmax [TFlop/s]</th><th>Rpeak [PFlop/s]</th><th>Power</th></tr><tr><td>{rank}</td><td><a href="/system/123/">Test</a><br>Lab<br>USA</td><td>10</td><td>1,500</td><td>2</td><td>1</td></tr></table>'
|
|
first = collector.parse_response(
|
|
html.format(rank=1), "https://top500.org/lists/top500/2025/11/"
|
|
)[0]
|
|
second = collector.parse_response(
|
|
html.format(rank=2), "https://top500.org/lists/top500/2026/06/"
|
|
)[0]
|
|
assert first["source_id"] == second["source_id"] == "top500_system_123"
|
|
assert float(first["value"]) == 1.5e15
|
|
assert first["metadata"]["rpeak_flops"] == 2e15
|
|
assert first["reference_date"] == "2025-11-01"
|
|
assert second["reference_date"] == "2026-06-01"
|
|
with pytest.raises(ValueError, match="date missing"):
|
|
collector.parse_response(html.format(rank=1))
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_large_source_loss_does_not_create_snapshot():
|
|
db = AsyncMock()
|
|
db.scalar.return_value = 500
|
|
db.add = MagicMock()
|
|
with pytest.raises(ValueError, match="40%"):
|
|
await TOP500Collector()._create_snapshot(db, 1, [{}] * 2, datetime.now(UTC))
|
|
db.add.assert_not_called()
|
|
db.commit.assert_not_awaited()
|
|
|
|
|
|
def test_history_missing_years_and_denominators():
|
|
global_data = compute_history(["CN", "US"])
|
|
assert global_data["compute"][3]["values"] == {"CN": 31, "US": 41}
|
|
assert global_data["compute"][-1]["values"] == {"CN": None, "US": None}
|
|
bilateral = compute_history(["CN", "US"], "bilateral_share")
|
|
assert sum(bilateral["compute"][3]["values"].values()) == 100
|
|
assert bilateral["denominator"] == "US+CN"
|
|
absolute = compute_history(["US", "CN"], "absolute")
|
|
assert absolute["compute"][3]["values"] == {"US": None, "CN": 435}
|
|
assert global_data["innovation"][-2]["value"] == 17.95
|
|
assert global_data["innovation"][-1]["value"] is None
|
|
assert global_data["evidence"]["innovation-2026"]["published_at"] == "2026-07-17"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_snapshot_query_uses_observation_cutoff_not_location():
|
|
query = await sample_query(AsyncMock(), ["CN"], 2023)
|
|
sql = str(query.compile(dialect=postgresql.dialect(), compile_kwargs={"literal_binds": True}))
|
|
assert "2024-01-01" in sql
|
|
assert "row_number() OVER (PARTITION BY" in sql
|
|
assert "compute_contract" in sql and "中国" in sql
|
|
assert "latitude" not in sql and "longitude" not in sql
|
|
|
|
|
|
def test_country_input_is_bounded_and_deduplicated():
|
|
assert selected_countries("CN,CN,US") == ["CN", "US"]
|
|
assert selected_countries("CN,US,EUROPE,APAC,OTHER")[-1] == "OTHER"
|
|
with pytest.raises(HTTPException):
|
|
selected_countries("invalid")
|
|
with pytest.raises(HTTPException, match="overlapping"):
|
|
selected_countries("EUROPE,DE")
|
|
|
|
|
|
def test_catalog_supports_global_participants_without_inventing_history():
|
|
from app.services.compute_catalog import (
|
|
CATALOG,
|
|
COMPARISON_GROUPS,
|
|
COUNTRIES,
|
|
source_country_names,
|
|
source_country_participants,
|
|
)
|
|
|
|
assert {"CN", "US", "JP", "KR", "IN", "DE", "FR", "GB", "AU", "BR", "ZA"}.issubset(
|
|
{entry["code"] for entry in CATALOG}
|
|
)
|
|
assert not {"TW", "HK", "MO", "XK"}.intersection(entry["code"] for entry in CATALOG)
|
|
owners = source_country_names(["CN"])
|
|
assert {"中国", "中国（台湾）", "中国（香港）", "中国（澳门）"}.issubset(owners)
|
|
assert source_country_participants(["CN"])["中国（台湾）"] == "CN"
|
|
assert [entry["code"] for entry in CATALOG[:5]] == ["CN", "US", "EUROPE", "APAC", "OTHER"]
|
|
assert [entry["kind"] for entry in CATALOG[:5]] == [
|
|
"country",
|
|
"country",
|
|
"region",
|
|
"region",
|
|
"region",
|
|
]
|
|
group_members = [
|
|
member
|
|
for _code, _name, _name_en, members, _anchor in COMPARISON_GROUPS
|
|
for member in members
|
|
]
|
|
assert len(group_members) == len(set(group_members))
|
|
assert set(group_members) == set(COUNTRIES)
|
|
data = compute_history(["JP", "KR"])
|
|
assert [row["values"]["JP"] for row in data["compute"][:4]] == [6, 5, 4, 4]
|
|
assert all(row["values"]["KR"] is None for row in data["compute"])
|
|
assert all(row["value"] is None for row in data["innovation"])
|
|
assert len(data["catalog"]) > 30
|
|
|
|
|
|
def test_region_history_requires_complete_same_year_member_coverage():
|
|
data = compute_history(["CN", "US", "EUROPE", "APAC", "OTHER"])
|
|
row_2023 = next(row for row in data["compute"] if row["year"] == 2023)
|
|
assert row_2023["values"]["CN"] == 31
|
|
assert row_2023["values"]["US"] == 41
|
|
assert row_2023["values"]["EUROPE"] is None
|
|
assert row_2023["values"]["APAC"] is None
|
|
assert row_2023["values"]["OTHER"] is None
|
|
assert row_2023["coverage"]["EUROPE"]["status"] == "incomplete"
|
|
assert "DE" in row_2023["coverage"]["EUROPE"]["missing_members"]
|
|
assert data["default_participants"] == ["CN", "US", "EUROPE", "APAC", "OTHER"]
|
|
|
|
|
|
def test_anonymous_records_never_enter_the_location_repair_queue():
|
|
from app.api.v1.visualization import convert_compute_centers_to_geojson
|
|
from app.models.collected_data import CollectedData
|
|
|
|
record = CollectedData(
|
|
source="epoch_ai_gpu",
|
|
source_id="anonymous",
|
|
name="Anonymous",
|
|
extra_data={"compute_contract": "v1", "anonymous": True},
|
|
)
|
|
result = convert_compute_centers_to_geojson([record])
|
|
assert result["features"] == []
|
|
assert result["unresolved"] == []
|
|
assert result["withheld_count"] == 1
|
|
|
|
|
|
def test_absolute_report_year_mapping_and_bilateral_denominator_ignore_other_countries():
|
|
absolute = compute_history(["CN"], "absolute")
|
|
observed = {row["year"]: row for row in absolute["compute"]}
|
|
assert observed[2022]["values"]["CN"] == 302
|
|
assert observed[2022]["evidence_id"] == "caict-2023"
|
|
assert observed[2023]["evidence_id"] == "caict-2024"
|
|
bilateral = compute_history(["CN", "US", "JP"], "bilateral_share")
|
|
values = next(row["values"] for row in bilateral["compute"] if row["year"] == 2023)
|
|
assert values == {"CN": 43.06, "US": 56.94, "JP": None}
|
|
for row in absolute["compute"]:
|
|
if row["evidence_id"]:
|
|
report = absolute["evidence"][row["evidence_id"]]
|
|
assert report["url"].startswith("https://www.caict.ac.cn/")
|