spb/company-atlas

1"""Every committed registry line is loadable and consistent; the universe is large and diversified (docs/SEEDS.md)."""2from __future__ import annotations34import csv5from collections import Counter67import pytest89from companyatlas.registry.industries import REGISTRY_DIR, is_valid_slug, load_industries, top_level_of, top_level_slugs10from companyatlas.registry.seed import load_registry_rows, registry_files11from companyatlas.urls import registrable_domain1213MIN_TOTAL = 600014US_CAP = 0.42            # 40 % target + tolerance for the industry top-up15OTHER_CAP = 0.1316REQUIRED = {"wikidata_id", "display_name", "website", "canonical_domain", "country", "industries", "importance", "tier", "source", "harvested_at"}171819@pytest.fixture(scope="module")20def rows() -> list[dict]:21    files = registry_files()22    assert files, "registry/companies/*.ndjson missing — run scripts/seed_wikidata.py"23    return load_registry_rows(files)242526@pytest.fixture(scope="module")27def countries() -> dict[str, dict]:28    with (REGISTRY_DIR / "countries.csv").open(encoding="utf-8") as f:29        return {r["code"]: r for r in csv.DictReader(f)}303132def test_countries_csv(countries: dict[str, dict]) -> None:33    assert len(countries) >= 24534    for code, r in countries.items():35        assert len(code) == 2 and code.isupper() and r["name"]36        assert -90 <= float(r["lat"]) <= 90 and -180 <= float(r["lon"]) <= 18037    for must in ("US", "CA", "GB", "DE", "FR", "JP", "KR", "IN", "AU", "BR", "NG", "TW", "HK", "XK"):38        assert must in countries39    assert countries["TW"]["region"] == "Asia"404142def test_every_line_is_valid(rows: list[dict], countries: dict[str, dict]) -> None:43    assert len(rows) >= MIN_TOTAL, f"only {len(rows)} companies"44    slugs = {i.slug for i in load_industries()}45    domains: Counter[str] = Counter()46    qids: Counter[str] = Counter()47    for r in rows:48        assert REQUIRED <= set(r), r.get("wikidata_id")49        assert r["wikidata_id"].startswith("Q") and r["wikidata_id"][1:].isdigit()50        assert r["display_name"].strip()51        assert r["website"].startswith(("https://", "http://")), r["website"]52        assert "/" not in r["website"].split("://", 1)[1], "website must be scheme + host only"53        assert r["canonical_domain"] == registrable_domain(r["website"])54        assert r["country"] is None or r["country"] in countries, r["country"]55        assert all(is_valid_slug(s) for s in r["industries"]), r["industries"]56        assert all(s in slugs for s in r["industries"])57        assert 0.0 <= r["importance"] <= 1.0 and r["tier"] in (1, 2, 3, 4)58        assert r["source"] == "wikidata" and r["sitelinks"] >= 259        assert r["founded_year"] is None or 1000 <= r["founded_year"] <= 202660        assert r["lat"] is None or -90 <= r["lat"] <= 9061        assert r["lon"] is None or -180 <= r["lon"] <= 18062        assert r["parent"] is None or set(r["parent"]) == {"wikidata_id", "name"}63        domains[r["canonical_domain"]] += 164        qids[r["wikidata_id"]] += 165    assert not [d for d, n in domains.items() if n > 1], "duplicate domains"66    assert not [q for q, n in qids.items() if n > 1], "duplicate wikidata ids"676869def test_diversification(rows: list[dict]) -> None:70    n = len(rows)71    by_country = Counter(r["country"] for r in rows)72    assert by_country["US"] / n <= US_CAP, by_country["US"] / n73    for code, count in by_country.items():74        if code not in ("US", None):75            assert count / n <= OTHER_CAP, (code, count / n)76    assert by_country[None] / n <= 0.0577    for code, minimum in {"CA": 150, "GB": 150, "DE": 150, "FR": 150, "JP": 150, "IN": 150, "AU": 150}.items():78        assert by_country[code] >= minimum, (code, by_country[code])79    tiers = Counter(r["tier"] for r in rows)80    assert tiers[1] == 150 and tiers[2] == 800 and tiers[3] == 2500 and tiers[4] == n - 345081    by_top: Counter[str] = Counter()82    for r in rows:83        for t in {top_level_of(s) for s in r["industries"]}:84            by_top[t] += 185    covered = [t for t in top_level_slugs() if by_top[t] >= 60]86    assert len(covered) >= len(top_level_slugs()) - 2, {t: by_top[t] for t in top_level_slugs() if by_top[t] < 60}87    assert sum(1 for r in rows if r["industries"]) / n >= 0.85888990def test_files_are_split_by_region(rows: list[dict], countries: dict[str, dict]) -> None:91    for path in registry_files():92        region = path.stem.removeprefix("wikidata-")93        for r in load_registry_rows([path]):94            actual = (countries.get(r["country"] or "", {}).get("region") or "other").lower().replace(" ", "-")95            assert actual == region, (path.name, r["wikidata_id"], actual)96