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 = 3000014US_CAP = 0.37            # 35 % target + tolerance for the industry top-up15OTHER_CAP = 0.1316TIER_SHARES = (0.01, 0.0667, 0.2667)     # quantiles of importance → tiers 1–3 (scripts/seed_wikidata.py TIER_SHARES)17REQUIRED = {"wikidata_id", "display_name", "website", "canonical_domain", "country", "industries", "importance", "tier", "source", "harvested_at"}181920def tier_cutoffs(n: int) -> tuple[int, int, int]:21    t1 = round(n * TIER_SHARES[0])22    t2 = t1 + round(n * TIER_SHARES[1])23    return t1, t2, t2 + round(n * TIER_SHARES[2])242526@pytest.fixture(scope="module")27def rows() -> list[dict]:28    files = registry_files()29    assert files, "registry/companies/*.ndjson missing — run scripts/seed_wikidata.py"30    return load_registry_rows(files)313233@pytest.fixture(scope="module")34def countries() -> dict[str, dict]:35    with (REGISTRY_DIR / "countries.csv").open(encoding="utf-8") as f:36        return {r["code"]: r for r in csv.DictReader(f)}373839def test_countries_csv(countries: dict[str, dict]) -> None:40    assert len(countries) >= 24541    for code, r in countries.items():42        assert len(code) == 2 and code.isupper() and r["name"]43        assert -90 <= float(r["lat"]) <= 90 and -180 <= float(r["lon"]) <= 18044    for must in ("US", "CA", "GB", "DE", "FR", "JP", "KR", "IN", "AU", "BR", "NG", "TW", "HK", "XK"):45        assert must in countries46    assert countries["TW"]["region"] == "Asia"474849def test_every_line_is_valid(rows: list[dict], countries: dict[str, dict]) -> None:50    assert len(rows) >= MIN_TOTAL, f"only {len(rows)} companies"51    slugs = {i.slug for i in load_industries()}52    domains: Counter[str] = Counter()53    qids: Counter[str] = Counter()54    for r in rows:55        assert REQUIRED <= set(r), r.get("wikidata_id")56        assert r["wikidata_id"].startswith("Q") and r["wikidata_id"][1:].isdigit()57        assert r["display_name"].strip()58        assert r["website"].startswith(("https://", "http://")), r["website"]59        assert "/" not in r["website"].split("://", 1)[1], "website must be scheme + host only"60        assert r["canonical_domain"] == registrable_domain(r["website"])61        assert r["country"] is None or r["country"] in countries, r["country"]62        assert all(is_valid_slug(s) for s in r["industries"]), r["industries"]63        assert all(s in slugs for s in r["industries"])64        assert 0.0 <= r["importance"] <= 1.0 and r["tier"] in (1, 2, 3, 4)65        # Listed companies and large companies (employees / revenue bands) may have no Wikipedia article at all; class-band items need ≥ 2.66        assert r["source"] == "wikidata" and isinstance(r["sitelinks"], int) and r["sitelinks"] >= 067        assert r["founded_year"] is None or 1000 <= r["founded_year"] <= 202668        assert r["lat"] is None or -90 <= r["lat"] <= 9069        assert r["lon"] is None or -180 <= r["lon"] <= 18070        assert r["parent"] is None or set(r["parent"]) == {"wikidata_id", "name"}71        domains[r["canonical_domain"]] += 172        qids[r["wikidata_id"]] += 173    assert not [d for d, n in domains.items() if n > 1], "duplicate domains"74    assert not [q for q, n in qids.items() if n > 1], "duplicate wikidata ids"757677def test_diversification(rows: list[dict]) -> None:78    n = len(rows)79    by_country = Counter(r["country"] for r in rows)80    assert by_country["US"] / n <= US_CAP, by_country["US"] / n81    for code, count in by_country.items():82        if code not in ("US", None):83            assert count / n <= OTHER_CAP, (code, count / n)84    assert by_country[None] / n <= 0.0585    for code, minimum in {"CA": 150, "GB": 150, "DE": 150, "FR": 150, "JP": 150, "IN": 150, "AU": 150}.items():86        assert by_country[code] >= minimum, (code, by_country[code])87    tiers = Counter(r["tier"] for r in rows)88    t1, t2, t3 = tier_cutoffs(n)89    assert (tiers[1], tiers[2], tiers[3], tiers[4]) == (t1, t2 - t1, t3 - t2, n - t3), dict(tiers)90    assert 250 <= tiers[1] <= 400 and 1800 <= tiers[2] <= 2500 and 7000 <= tiers[3] <= 1000091    # The atlas is listed-companies-first: a substantial share carries a stock exchange or ticker.92    assert sum(1 for r in rows if r["public_company"]) / n >= 0.3093    by_top: Counter[str] = Counter()94    for r in rows:95        for t in {top_level_of(s) for s in r["industries"]}:96            by_top[t] += 197    covered = [t for t in top_level_slugs() if by_top[t] >= 60]98    assert len(covered) >= len(top_level_slugs()) - 2, {t: by_top[t] for t in top_level_slugs() if by_top[t] < 60}99    # Listed small caps (TSE, Bursa Malaysia, KRX, TASE …) often carry no P452 and a bare "public company" description on Wikidata; the100    # LLM classify_industry job and the crawl fill those later (docs/SEEDS.md, Known gaps).101    assert sum(1 for r in rows if r["industries"]) / n >= 0.80102103104def test_files_are_split_by_region(rows: list[dict], countries: dict[str, dict]) -> None:105    for path in registry_files():106        region = path.stem.removeprefix("wikidata-")107        for r in load_registry_rows([path]):108            actual = (countries.get(r["country"] or "", {}).get("region") or "other").lower().replace(" ", "-")109            assert actual == region, (path.name, r["wikidata_id"], actual)110