spb/company-atlas

1"""Test data factories for the intelligence layer (slug prefix `ztest-`; `cleanup()` removes everything by cascade)."""2from __future__ import annotations34import contextlib5import uuid6from datetime import UTC, datetime, timedelta7from typing import Any89import pytest1011from companyatlas.db import dispose, execute, fetch_one, jsonb, transaction12from companyatlas.ids import new_id, normalize_alias, stable_hash1314PREFIX = "ztest-"15CONNECTOR_HTML = "ztest-generic-html-v1"16CONNECTOR_ATS = "ztest-greenhouse-v1"171819def _uid() -> str:20    return uuid.uuid4().hex[:8]212223async def ensure_reference(conn) -> None:  # type: ignore[no-untyped-def]24    for code, name in (("CA", "Canada"), ("US", "United States"), ("JP", "Japan"), ("DE", "Germany"), ("GB", "United Kingdom")):25        await execute(conn, "insert into countries (code, name) values (:c, :n) on conflict (code) do nothing", c=code, n=name)26    for cid, cat, mode in ((CONNECTOR_HTML, "homepage", "http"), (CONNECTOR_ATS, "jobs_board", "json")):27        await execute(conn, "insert into connectors (id, name, version, category, fetch_mode) values (:id, :id, 'v1', :cat, :mode) on conflict (id) do nothing",28                      id=cid, cat=cat, mode=mode)293031async def make_company(conn, *, name: str | None = None, country: str = "CA", industries: list[str] | None = None,  # type: ignore[no-untyped-def]32                       first_observed_days_ago: int | None = 30) -> dict[str, Any]:33    await ensure_reference(conn)34    uid = _uid()35    slug = f"{PREFIX}{uid}"36    display = name or f"ZTest Corp {uid}"37    cid = new_id("company")38    first = datetime.now(UTC) - timedelta(days=first_observed_days_ago) if first_observed_days_ago is not None else None39    await execute(conn, """40        insert into companies (id, slug, display_name, canonical_domain, website, industries, country, onboarding_status, first_observed_at, last_observed_at)41        values (:id, :slug, :name, :domain, :web, cast(:ind as text[]), :country, 'active', :first, :first)""",42        id=cid, slug=slug, name=display, domain=f"{slug}.example", web=f"https://{slug}.example", ind=industries or [], country=country, first=first)43    await execute(conn, "insert into company_aliases (company_id, alias, alias_norm) values (:c, :a, :n) on conflict do nothing", c=cid, a=display, n=normalize_alias(display))44    return {"id": cid, "slug": slug, "display_name": display, "country": country, "industries": industries or [], "canonical_domain": f"{slug}.example"}454647async def make_sensor(conn, company: dict[str, Any], surface: str, *, connector_id: str = CONNECTOR_HTML, path: str | None = None,  # type: ignore[no-untyped-def]48                      status: str = "active", created_days_ago: int = 30, interval_s: int = 86400) -> dict[str, Any]:49    sid = new_id("sensor")50    url = f"https://{company['canonical_domain']}/{path or surface}"51    await execute(conn, """52        insert into sensors (id, company_id, surface, connector_id, url, canonical_url, domain, status, base_interval_s, current_interval_s, created_at, last_success_at)53        values (:id, :c, :surface, :conn, :url, :url, :domain, :status, :iv, :iv, :created, now())""",54        id=sid, c=company["id"], surface=surface, conn=connector_id, url=url, domain=company["canonical_domain"], status=status, iv=interval_s,55        created=datetime.now(UTC) - timedelta(days=created_days_ago))56    return {"id": sid, "company_id": company["id"], "surface": surface, "connector_id": connector_id, "url": url}575859async def make_snapshot(conn, sensor: dict[str, Any], *, version_no: int = 1, title: str | None = None, fetched_at: datetime | None = None) -> str:  # type: ignore[no-untyped-def]60    snap = new_id("snapshot")61    h = stable_hash(snap)62    await execute(conn, """63        insert into snapshots (id, sensor_id, company_id, version_no, fetched_at, content_hash, normalized_hash, structural_hash, title)64        values (:id, :s, :c, :v, :at, :h, :h, :h, :title)""", id=snap, s=sensor["id"], c=sensor["company_id"], v=version_no, at=fetched_at or datetime.now(UTC), h=h, title=title)65    return snap666768async def make_change(conn, sensor: dict[str, Any], *, significance: float = 0.6, kind: str | None = None, structured_delta: dict[str, Any] | None = None,  # type: ignore[no-untyped-def]69                      diff: dict[str, Any] | None = None, detected_at: datetime | None = None, status: str = "pending") -> dict[str, Any]:70    from companyatlas.taxonomy import change_kind7172    kind = kind or str(change_kind(significance))73    before = await make_snapshot(conn, sensor, version_no=1, fetched_at=(detected_at or datetime.now(UTC)) - timedelta(days=1))74    after = await make_snapshot(conn, sensor, version_no=2, fetched_at=detected_at)75    diff = diff or {"added": [], "removed": [], "modified": [], "moved": [], "counts": {"added": 0, "removed": 0, "modified": 0, "moved": 0},76                    "text_delta_ratio": 0.1, "similarity": 0.9, "reasons": []}77    cid = new_id("change")78    await execute(conn, """79        insert into changes (id, sensor_id, company_id, surface, snapshot_before, snapshot_after, detected_at, significance, kind, blocks_added, blocks_removed,80                             blocks_modified, text_delta_ratio, similarity, diff, structured_delta, status)81        values (:id, :s, :c, :surface, :before, :after, :at, :sig, :kind, :ba, :br, :bm, :ratio, :sim, cast(:diff as jsonb), cast(:delta as jsonb), :status)""",82        id=cid, s=sensor["id"], c=sensor["company_id"], surface=sensor["surface"], before=before, after=after, at=detected_at or datetime.now(UTC), sig=significance,83        kind=kind, ba=len(diff.get("added") or []), br=len(diff.get("removed") or []), bm=len(diff.get("modified") or []), ratio=diff.get("text_delta_ratio") or 0,84        sim=diff.get("similarity"), diff=jsonb(diff), delta=jsonb(structured_delta or {}), status=status)85    return await fetch_one(conn, "select * from changes where id = :id", id=cid) or {"id": cid}868788async def make_job(conn, company: dict[str, Any], *, title: str = "Software Engineer", first_seen_days_ago: float = 10, removed_days_ago: float | None = None,  # type: ignore[no-untyped-def]89                   country: str | None = "CA", is_ai: bool = False, remote: bool = False, department: str | None = "Engineering", sensor_id: str | None = None) -> str:90    jid = new_id("job")91    now = datetime.now(UTC)92    first = now - timedelta(days=first_seen_days_ago)93    removed = now - timedelta(days=removed_days_ago) if removed_days_ago is not None else None94    await execute(conn, """95        insert into jobs (id, company_id, sensor_id, fingerprint, title, department, country, remote, first_seen_at, last_seen_at, removed_at, status, is_ai)96        values (:id, :c, :s, :fp, :title, :dep, :country, :remote, :first, :last, :removed, :status, :ai)""",97        id=jid, c=company["id"], s=sensor_id, fp=stable_hash(jid), title=title, dep=department, country=country, remote=remote, first=first,98        last=removed or now, removed=removed, status="no_longer_listed" if removed else "open", ai=is_ai)99    return jid100101102async def make_event(conn, company: dict[str, Any], *, subtype: str, importance: float = 0.6, days_ago: float = 1, title: str | None = None,  # type: ignore[no-untyped-def]103                     tags: list[str] | None = None, surface: str | None = None, sensor_id: str | None = None) -> str:104    from companyatlas.taxonomy import EVENT_SUBTYPES, EventType, confidence_label105106    eid = new_id("event")107    etype = str(EVENT_SUBTYPES.get(subtype, (EventType.OTHER, 0.3))[0])108    await execute(conn, """109        insert into events (id, company_id, sensor_id, surface, event_type, event_subtype, importance, confidence, confidence_label, title, tags, detected_at, dedupe_key)110        values (:id, :c, :s, :surface, :t, :st, :imp, 0.8, :label, :title, cast(:tags as text[]), :at, :dk)""",111        id=eid, c=company["id"], s=sensor_id, surface=surface, t=etype, st=subtype, imp=importance, label=confidence_label(0.8),112        title=title or f"{subtype} test event", tags=tags or [], at=datetime.now(UTC) - timedelta(days=days_ago), dk=stable_hash(eid))113    return eid114115116async def make_location(conn, company: dict[str, Any], *, name: str, country: str, city: str | None = None, first_seen_days_ago: float = 5) -> str:  # type: ignore[no-untyped-def]117    lid = new_id("location")118    await execute(conn, """insert into locations (id, company_id, kind, name, name_norm, city, country, first_seen_at, last_seen_at)119                           values (:id, :c, 'office', :name, :norm, :city, :country, :first, now())""",120                  id=lid, c=company["id"], name=name, norm=normalize_alias(name) + _uid(), city=city, country=country, first=datetime.now(UTC) - timedelta(days=first_seen_days_ago))121    return lid122123124async def make_observation(conn, sensor: dict[str, Any], *, days_ago: float, changed: bool = False, not_modified: bool = False) -> str:  # type: ignore[no-untyped-def]125    oid = new_id("observation")126    await execute(conn, """insert into observations (id, sensor_id, company_id, fetched_at, status_code, changed, not_modified)127                           values (:id, :s, :c, :at, 200, :changed, :nm)""", id=oid, s=sensor["id"], c=sensor["company_id"],128                  at=datetime.now(UTC) - timedelta(days=days_ago), changed=changed, nm=not_modified)129    return oid130131132@pytest.fixture133async def intel_db():134    """DB fixture for the intelligence tests: a fresh engine bound to this test's event loop (other modules run on a session-scoped loop),135    skip when Postgres is unreachable, dispose afterwards. Import it into a test module to register it."""136    from companyatlas.config import settings137    from companyatlas.db import fetch_val138139    with contextlib.suppress(Exception):140        await dispose()141    try:142        async with transaction() as conn:143            ok = (await fetch_val(conn, "select 1")) == 1144    except Exception:  # noqa: BLE001145        ok = False146        with contextlib.suppress(Exception):147            await dispose()148    if not ok:149        pytest.skip(f"database not reachable: {settings.database_url.split('@')[-1]}")150    try:151        yield152    finally:153        with contextlib.suppress(Exception):154            await dispose()155156157async def cleanup() -> None:158    async with transaction() as conn:159        await execute(conn, "delete from companies where slug like :p", p=f"{PREFIX}%")160        await execute(conn, "delete from connectors where id like :p", p=f"{PREFIX}%")161        await execute(conn, "delete from alerts where name like :p", p=f"{PREFIX}%")162        await execute(conn, "delete from owners where token_hash like :p", p=f"{PREFIX}%")163        await execute(conn, "delete from global_daily where day < '2002-01-01'")164        await execute(conn, "delete from trends where day < '2002-01-01'")165166167__all__ = ["CONNECTOR_ATS", "CONNECTOR_HTML", "PREFIX", "cleanup", "ensure_reference", "intel_db", "make_change", "make_company", "make_event", "make_job", "make_location",168           "make_observation", "make_sensor", "make_snapshot"]169