spb/company-atlas

1"""Integration: the pipeline against the local Postgres with a temporary `ztest-` company (skipped when the DB is unavailable)."""2from __future__ import annotations34import asyncio5from typing import Any67import pytest8from conftest import FakeFetcher, fixture_path910from companyatlas.fetch import FetchError, NotModified, file_result11from companyatlas.taxonomy import FailureClass1213BASE = "https://www.acme-cloud.example"141516async def _setup() -> tuple[str, dict[str, str]]:17    from companyatlas.db import execute, jsonb, transaction18    from companyatlas.ids import new_id1920    cid = new_id("company")21    slug = f"ztest-pipe-{cid[-6:].lower()}"22    dom = f"{slug}.example"23    sensors: dict[str, str] = {}24    async with transaction() as conn:25        await execute(conn, """insert into companies (id, slug, display_name, canonical_domain, website, tier, importance, onboarding_status)26                               values (:id, :slug, 'Pipeline test', :dom, :web, 3, 0.4, 'active')""", id=cid, slug=slug, dom=dom, web=f"https://{dom}/")27        for surface, url, con in [("pricing", BASE + "/pricing", "generic-html-v1"), ("careers", BASE + "/careers", "generic-html-v1"),28                                  ("jobs_board", "https://boards-api.greenhouse.io/v1/boards/ztest/jobs?content=false", "greenhouse-v1"),29                                  ("leadership", BASE + "/about/leadership", "generic-html-v1"), ("locations", BASE + "/company/locations", "generic-html-v1"),30                                  ("newsroom", BASE + "/news", "generic-html-v1"), ("legal_terms", BASE + "/legal/terms", "generic-html-v1")]:31            sid = new_id("sensor")32            await execute(conn, """insert into sensors (id, company_id, surface, connector_id, url, canonical_url, domain, base_interval_s, current_interval_s, status, config)33                                   values (:id, :c, :s, :con, :u, :u, :d, 43200, 43200, 'pending', cast(:cfg as jsonb))""",34                          id=sid, c=cid, s=surface, con=con, u=url, d=dom, cfg=jsonb({"canonical_domain": "acme-cloud.example", "token": "ztest"}))35            sensors[surface] = sid36    return cid, sensors373839async def _teardown(cid: str) -> None:40    from companyatlas.db import dispose, execute, transaction4142    async with transaction() as conn:43        await execute(conn, "delete from companies where id = :c", c=cid)44        await execute(conn, "delete from domain_budgets where domain like 'ztest-%'")45    await dispose()464748async def _sensor(sid: str) -> dict[str, Any]:49    from companyatlas.db import fetch_one, transaction5051    async with transaction() as conn:52        row = await fetch_one(conn, "select * from sensors where id = :id", id=sid)53    assert row is not None54    return row555657async def _run(sid: str, path: str, ct: str = "text/html", fetcher: Any = None):  # type: ignore[no-untyped-def]58    from companyatlas.services.pipeline import run_sensor5960    row = await _sensor(sid)61    return await run_sensor(row, fetcher=fetcher or FakeFetcher(), worker="pytest", result=file_result(path, url=row["url"], content_type=ct))626364async def _q(sql: str, **params: Any) -> list[dict[str, Any]]:65    from companyatlas.db import fetch_all, transaction6667    async with transaction() as conn:68        return await fetch_all(conn, sql, **params)697071@pytest.mark.usefixtures("db")72def test_pipeline_end_to_end() -> None:73    async def scenario() -> None:74        cid, s = await _setup()75        try:76            gh = fixture_path("greenhouse", "stripe_jobs.json")77            # ---- baseline runs (version 1, no change rows, entities inserted)78            o = await _run(s["pricing"], fixture_path("generic_html", "pricing.html"))79            assert o.status == "ok" and o.snapshot_id and not o.change_id and o.delta_counts == {"plans_added": 3}80            # ---- same content again → unchanged, no snapshot, interval grows ×1.2581            o = await _run(s["pricing"], fixture_path("generic_html", "pricing.html"))82            assert o.status == "unchanged" and o.snapshot_id is None and o.interval_s == int(43200 * 1.25)83            # ---- price change → new snapshot version, change row pending, plan superseded, burst interval84            o = await _run(s["pricing"], fixture_path("generic_html", "pricing_v2.html"))85            assert o.status == "changed" and o.kind in ("major", "critical") and o.change_id and o.delta_counts.get("plans_price_changed") == 186            assert o.interval_s == 90087            plans = await _q("select plan_name, price, status, version_no from pricing_plans where company_id = :c order by plan_name, version_no", c=cid)88            starter = [p for p in plans if p["plan_name"] == "Starter"]89            assert [(int(p["price"]), p["status"], p["version_no"]) for p in starter] == [(29, "superseded", 1), (39, "current", 2)]90            chg = (await _q("select * from changes where company_id = :c and surface = 'pricing'", c=cid))[0]91            assert chg["status"] == "pending" and chg["structured_delta"]["plans"]["price_changed"][0]["after"] == 39 and chg["diff"]["counts"]["modified"] == 192            snaps = await _q("select version_no, previous_snapshot_id from snapshots where sensor_id = :s order by version_no", s=s["pricing"])93            assert [x["version_no"] for x in snaps] == [1, 2] and snaps[1]["previous_snapshot_id"] == (await _q("select id from snapshots where sensor_id = :s and version_no = 1", s=s["pricing"]))[0]["id"]94            # ---- careers: jobs added / removed with no_longer_listed95            await _run(s["careers"], fixture_path("generic_html", "careers.html"))96            o = await _run(s["careers"], fixture_path("generic_html", "careers_v2.html"))97            assert o.delta_counts == {"jobs_added": 2, "jobs_removed": 1}98            jobs = await _q("select title, status, removed_at, is_ai from jobs where sensor_id = :s order by title", s=s["careers"])99            by = {j["title"]: j for j in jobs}100            assert by["People Operations Lead"]["status"] == "no_longer_listed" and by["People Operations Lead"]["removed_at"] is not None101            assert by["AI Research Scientist"]["is_ai"] and by["AI Research Scientist"]["status"] == "open"102            assert len(jobs) == 7103            # ---- untrusted extraction never mass-removes: a page with no job listing leaves the 6 open jobs alone104            o = await _run(s["careers"], fixture_path("generic_html", "legal_terms.html"))105            open_after = await _q("select count(*) n from jobs where sensor_id = :s and status = 'open'", s=s["careers"])106            assert open_after[0]["n"] == 6107            # ---- structured board: baseline + unchanged108            o = await _run(s["jobs_board"], gh, "application/json")109            assert o.status == "ok" and o.delta_counts == {"jobs_added": 20}110            o = await _run(s["jobs_board"], gh, "application/json")111            assert o.status == "unchanged"112            # ---- other surfaces baseline113            for surf, fn in (("leadership", "leadership.html"), ("locations", "locations.html"), ("newsroom", "newsroom.html"), ("legal_terms", "legal_terms.html")):114                o = await _run(s[surf], fixture_path("generic_html", fn))115                assert o.status == "ok", (surf, o.error)116            assert (await _q("select count(*) n from people where company_id = :c and is_executive", c=cid))[0]["n"] == 4117            assert (await _q("select count(*) n from locations where company_id = :c and country = 'US'", c=cid))[0]["n"] == 2118            assert (await _q("select count(*) n from news_items where company_id = :c", c=cid))[0]["n"] == 4119            # ---- failure path: fetch error → observation + failures row + backoff; repeated → failing120            from companyatlas.services.pipeline import run_sensor121122            ff = FakeFetcher()123            ff.add_error(BASE + "/legal/terms", FetchError("http 404", status=404, url=BASE + "/legal/terms", failure=FailureClass.PAGE_REMOVED))124            row = await _sensor(s["legal_terms"])125            o = await run_sensor(row, fetcher=ff, worker="pytest")  # type: ignore[arg-type]126            assert o.status == "failed" and o.failure_class == "PAGE_REMOVED" and o.interval_s == min(7 * 86400, int(row["current_interval_s"] * 4.0))127            o = await run_sensor(await _sensor(s["legal_terms"]), fetcher=ff, worker="pytest")  # type: ignore[arg-type]128            assert o.sensor_status == "failing"129            fails = await _q("select failure_class from failures where sensor_id = :s", s=s["legal_terms"])130            assert len(fails) == 2131            # ---- not modified path132            ff2 = FakeFetcher()133            ff2.add_error(BASE + "/news", NotModified(5))134            o = await run_sensor(await _sensor(s["newsroom"]), fetcher=ff2, worker="pytest")  # type: ignore[arg-type]135            assert o.status == "not_modified"136            obs = await _q("select not_modified from observations where sensor_id = :s order by fetched_at desc limit 1", s=s["newsroom"])137            assert obs[0]["not_modified"] is True138            # ---- robots block → review queue item139            ff3 = FakeFetcher()140            from companyatlas.fetch import BlockedError141142            ff3.add_error(BASE + "/company/locations", BlockedError("robots", url=BASE + "/company/locations", failure=FailureClass.ROBOTS))143            o = await run_sensor(await _sensor(s["locations"]), fetcher=ff3, worker="pytest")  # type: ignore[arg-type]144            assert o.sensor_status == "blocked"145            assert (await _q("select count(*) n from review_queue where ref_id = :s and kind = 'blocked_source'", s=s["locations"]))[0]["n"] == 1146            # ---- company + ledger touched147            comp = (await _q("select stats, first_observed_at, last_change_at from companies where id = :c", c=cid))[0]148            assert comp["first_observed_at"] is not None and comp["last_change_at"] is not None and comp["stats"]["meaningful_changes"] >= 2149            led = await _q("select units from cost_ledger where day = current_date and dimension = 'fetch' and key = :c", c=cid)150            assert led and led[0]["units"] >= 10151        finally:152            await _teardown(cid)153154    asyncio.run(scenario())155156157@pytest.mark.usefixtures("db")158def test_scheduler_claims_and_releases() -> None:159    async def scenario() -> None:160        from companyatlas.services.scheduler import claim_due_sensors, release_claims161162        cid, _sensors = await _setup()163        try:164            rows = await claim_due_sensors("pytest-worker", 100)165            mine = [r for r in rows if r["company_id"] == cid]166            assert len(mine) == 7 and all(r["claimed_by"] == "pytest-worker" for r in mine)167            again = await claim_due_sensors("pytest-worker-2", 100)168            assert not [r for r in again if r["company_id"] == cid]          # already claimed169            await release_claims("pytest-worker")170            assert (await _q("select count(*) n from sensors where company_id = :c and claimed_by is null", c=cid))[0]["n"] == 7171        finally:172            await _teardown(cid)173174    asyncio.run(scenario())175176177@pytest.mark.usefixtures("db")178def test_connectors_table_sync() -> None:179    async def scenario() -> None:180        from companyatlas.db import dispose, transaction181        from companyatlas.sdk.connector import all_connectors, sync_connectors_table182183        async with transaction() as conn:184            n = await sync_connectors_table(conn)185        assert n == len(all_connectors()) >= 14186        rows = await _q("select id from connectors where id = 'generic-html-v1'")187        assert rows188        await dispose()189190    asyncio.run(scenario())191