spb/company-atlas

1"""Seed loader against the local database: idempotent upserts, aliases, relationships, discover queue (rows use `ztest-` slugs)."""2from __future__ import annotations34import json5from pathlib import Path67import pytest89from companyatlas.db import dispose, fetch_all, fetch_one, fetch_val, transaction10from companyatlas.registry.seed import add_company, normalise_row, seed, unique_slug1112FIXTURE = [13    {"wikidata_id": "Q900000001", "display_name": "ZTest Alpha Corp", "legal_name": "ZTest Alpha Corporation", "aliases": ["ZTest Alpha", "Alpha Systems"],14     "website": "https://www.ztest-alpha.com/", "canonical_domain": "ztest-alpha.com", "country": "CA", "hq_city": "Montréal",15     "industries": ["software", "cloud-infrastructure"], "industry_labels": ["software industry"], "founded_year": 1999, "employees": 1200,16     "public_company": True, "ticker": "ZTA", "exchange": "Toronto Stock Exchange", "parent": None, "sitelinks": 40, "importance": 0.8, "tier": 1,17     "source": "wikidata", "harvested_at": "2026-09-12T00:00:00+00:00"},18    {"wikidata_id": "Q900000002", "display_name": "ZTest Beta", "website": "https://ztest-beta.com", "country": "US", "industries": ["fintech"],19     "parent": {"wikidata_id": "Q900000001", "name": "ZTest Alpha Corp"}, "importance": 0.4, "tier": 3, "source": "wikidata"},20    {"website": "ztest-gamma.com", "country": "QZ"},                                        # no name, unknown country → derived / nulled21    {"wikidata_id": "Q900000004", "display_name": "ZTest Alpha Duplicate", "website": "https://shop.ztest-alpha.com/x"},   # same domain → skipped22]232425async def _cleanup() -> None:26    async with transaction() as conn:27        ids = [r["id"] for r in await fetch_all(conn, "select id from companies where slug like 'ztest-%'")]28        if ids:29            await conn.exec_driver_sql("delete from queue_jobs where key = any($1::text[])", ([f"discover:{i}" for i in ids],))30            await conn.exec_driver_sql("delete from companies where id = any($1::text[])", (ids,))313233@pytest.fixture34def fixture_file(tmp_path: Path) -> Path:35    p = tmp_path / "ztest.ndjson"36    p.write_text("\n".join(json.dumps(r) for r in FIXTURE) + "\n", encoding="utf-8")37    return p383940async def test_seed_twice_is_idempotent(fixture_file: Path) -> None:41    await _cleanup()42    try:43        async with transaction() as conn:44            first = await seed(conn, files=[fixture_file])45        assert first["industries"] >= 45 and first["countries"] >= 24046        assert first["companies_seen"] == 4 and first["companies_new"] == 3 and first["companies_skipped"] == 147        assert first["relationships"] == 2 and first["queue_jobs"] == 34849        async with transaction() as conn:50            rows = await fetch_all(conn, "select * from companies where slug like 'ztest-%' order by slug")51            assert [r["slug"] for r in rows] == ["ztest-alpha-corp", "ztest-beta", "ztest-gamma"]52            alpha, beta, gamma = rows53            assert alpha["onboarding_status"] == "pending" and alpha["country"] == "CA" and alpha["industry_primary"] == "software"54            assert alpha["public_company"] is True and alpha["ticker"] == "ZTA" and alpha["tier"] == 1 and abs(alpha["importance"] - 0.8) < 1e-655            assert alpha["source_meta"]["source"] == "wikidata" and alpha["source_meta"]["sitelinks"] == 4056            assert gamma["display_name"] == "Ztest Gamma" and gamma["country"] is None and gamma["source_meta"]["unknown_country"] == "QZ"57            aliases = await fetch_all(conn, "select alias_norm, kind from company_aliases where company_id = :id order by alias_norm", id=alpha["id"])58            # "ZTest Alpha Corp", "ZTest Alpha Corporation" and "ZTest Alpha" collapse to one key (legal suffixes dropped by normalize_alias)59            assert {a["alias_norm"] for a in aliases} == {"ztestalpha", "zta", "alphasystems"}60            assert {a["kind"] for a in aliases} >= {"brand", "ticker"}61            assert await fetch_val(conn, "select count(*) from domains where company_id = :id and kind = 'primary'", id=alpha["id"]) == 162            rels = await fetch_all(conn, "select from_company_id, to_company_id, kind, confidence from company_relationships "63                                         "where from_company_id in (:a, :b) order by kind", a=alpha["id"], b=beta["id"])64            assert [(r["kind"], r["from_company_id"] == alpha["id"]) for r in rels] == [("PARENT_OF", True), ("SUBSIDIARY_OF", False)]65            assert all(abs(r["confidence"] - 0.8) < 1e-6 for r in rels)66            job = await fetch_one(conn, "select kind, status, priority, payload from queue_jobs where key = :k", k=f"discover:{alpha['id']}")67            assert job and job["kind"] == "discover" and job["status"] == "pending" and abs(job["priority"] - 0.8) < 1e-668            assert job["payload"]["company_id"] == alpha["id"]69            assert (await fetch_one(conn, "select value from settings_kv where key = 'seed:last_run'"))["value"]["counters"]["companies_new"] == 37071        # Second run: nothing new, non-null values untouched even if the file changed, importance refreshed.72        changed = [dict(r) for r in FIXTURE]73        changed[0]["display_name"] = "RENAMED"74        changed[0]["legal_name"] = "Other Legal"75        changed[0]["importance"] = 0.576        changed[0]["hq_region"] = "Quebec"      # was null → filled77        fixture_file.write_text("\n".join(json.dumps(r) for r in changed) + "\n", encoding="utf-8")78        async with transaction() as conn:79            second = await seed(conn, files=[fixture_file])80        assert second["companies_new"] == 0 and second["companies_updated"] == 3 and second["queue_jobs"] == 0 and second["relationships"] == 081        async with transaction() as conn:82            alpha2 = await fetch_one(conn, "select * from companies where wikidata_id = 'Q900000001'")83            assert alpha2["display_name"] == "ZTest Alpha Corp" and alpha2["legal_name"] == "ZTest Alpha Corporation"84            assert alpha2["hq_region"] == "Quebec" and abs(alpha2["importance"] - 0.5) < 1e-685            assert await fetch_val(conn, "select count(*) from companies where slug like 'ztest-%'") == 386            assert await fetch_val(conn, "select count(*) from company_relationships r join companies c on c.id = r.from_company_id "87                                         "where c.slug like 'ztest-%'") == 288            assert await fetch_val(conn, "select count(*) from queue_jobs where key like 'discover:%' and payload->>'company_id' in "89                                         "(select id from companies where slug like 'ztest-%')") == 39091        # add_company on an existing domain tops up; a new one is created pending with a job.92        async with transaction() as conn:93            res = await add_company(conn, "ztest-alpha.com", display_name="ignored")94            assert res["slug"] == "ztest-alpha-corp" and res["counters"]["companies_new"] == 095            res = await add_company(conn, "https://ztest-delta.com/about", display_name="ZTest Delta", country="FR", industries=["banking"])96            assert res["slug"] == "ztest-delta" and res["onboarding_status"] == "pending" and res["counters"]["queue_jobs"] == 197    finally:98        await _cleanup()99        await dispose()100101102def test_normalise_and_slug_helpers() -> None:103    assert normalise_row({"website": ""}, source="x") is None104    assert normalise_row({"website": "ftp://nope"}, source="x") is None105    row = normalise_row({"website": "example.org", "industries": ["software", "not-a-slug"], "industry": "banking", "tier": 9, "importance": 3},106                        source="manual")107    assert row is not None108    assert row["website"] == "https://example.org" and row["canonical_domain"] == "example.org" and row["display_name"] == "Example"109    assert row["industries"] == ["banking", "software"] and row["tier"] == 4 and row["importance"] == 1.0110    assert unique_slug("acme", "US", set()) == "acme"111    assert unique_slug("acme", "US", {"acme"}) == "acme-us"112    assert unique_slug("acme", "US", {"acme", "acme-us"}) == "acme-2"113    assert unique_slug("acme", None, {"acme", "acme-2"}) == "acme-3"114