spb/company-atlas

1"""Shared fixtures for the API tests: a `ztest-api-*` company family with sensors, snapshots (real objects in the archive), a change,2events, jobs, entities, metrics, a signal, queue/failure/review rows and a paid API key. Everything is removed in the finaliser.34Import into a test module with `from test_api_support import client, fixture_data  # noqa: F401` and mark the module with5`pytestmark = pytest.mark.asyncio(loop_scope="session")` so the shared asyncpg engine stays on one loop.6"""7from __future__ import annotations89import hashlib10import json11import os12import secrets13from datetime import UTC, datetime, timedelta14from typing import Any1516import httpx17import pytest_asyncio1819from companyatlas import archive20from companyatlas.api.common import cache21from companyatlas.db import dispose, execute, fetch_val, jsonb, transaction2223ADMIN = {"X-CA-Admin-Token": "dev-admin-token"}24OWNER_TOKEN = "ztest-owner-token-" + secrets.token_hex(12)25OWNER = {"X-CA-Owner-Token": OWNER_TOKEN}26RAW_API_KEY = "ca_paid_ztest_" + secrets.token_urlsafe(24)27SUFFIX = secrets.token_hex(3)282930def _now() -> datetime:31    return datetime.now(UTC)323334async def _insert_fixture() -> dict[str, Any]:35    now = _now()36    d: dict[str, Any] = {"suffix": SUFFIX, "objects": []}37    ids = {"alpha": f"co_ztestapi{SUFFIX}a", "beta": f"co_ztestapi{SUFFIX}b", "sensor_home": f"sen_ztestapi{SUFFIX}h", "sensor_careers": f"sen_ztestapi{SUFFIX}c",38           "sensor_beta": f"sen_ztestapi{SUFFIX}x", "snap1": f"snap_ztestapi{SUFFIX}1", "snap2": f"snap_ztestapi{SUFFIX}2", "snap3": f"snap_ztestapi{SUFFIX}3",39           "change": f"chg_ztestapi{SUFFIX}1", "ev_hiring": f"evt_ztestapi{SUFFIX}h", "ev_pricing": f"evt_ztestapi{SUFFIX}p", "ev_product": f"evt_ztestapi{SUFFIX}n",40           "ev_leader": f"evt_ztestapi{SUFFIX}l", "ev_retracted": f"evt_ztestapi{SUFFIX}r", "queue_dead": f"qj_ztestapi{SUFFIX}d", "failure": f"fail_ztestapi{SUFFIX}1",41           "review": f"rev_ztestapi{SUFFIX}1", "signal": f"sig_ztestapi{SUFFIX}1", "api_key": f"key_ztestapi{SUFFIX}1", "connector": f"ztest-generic-{SUFFIX}"}42    d.update(ids)43    d["alpha_slug"], d["beta_slug"] = f"ztest-api-alpha-{SUFFIX}", f"ztest-api-beta-{SUFFIX}"44    d["country"], d["industry"] = "ZZ", f"ztest-industry-{SUFFIX}"4546    text1 = "Ztest Alpha — pricing\nStarter plan $10 per month\nPro plan $20 per month\nContact sales for Enterprise"47    text2 = "Ztest Alpha — pricing\nStarter plan $12 per month\nPro plan $20 per month\nContact sales for Enterprise\nNew: Team plan"48    text3 = text2 + "\nFooter updated"49    blocks1 = [{"key": "b1", "kind": "heading", "text": "Ztest Alpha — pricing", "path": "", "hash": "h1", "simhash": 0, "weight": 1.0, "order": 0, "attrs": {}},50               {"key": "b2", "kind": "pricing_plan", "text": "Starter plan $10 per month", "path": "Pricing > Starter", "hash": "h2", "simhash": 0, "weight": 1.5, "order": 1, "attrs": {}},51               {"key": "b3", "kind": "pricing_plan", "text": "Pro plan $20 per month", "path": "Pricing > Pro", "hash": "h3", "simhash": 0, "weight": 1.5, "order": 2, "attrs": {}}]52    blocks2 = [dict(blocks1[0]), {**blocks1[1], "text": "Starter plan $12 per month", "hash": "h2b"}, dict(blocks1[2]),53               {"key": "b4", "kind": "pricing_plan", "text": "New: Team plan", "path": "Pricing > Team", "hash": "h4", "simhash": 0, "weight": 1.5, "order": 3, "attrs": {}}]54    blocks3 = blocks2 + [{"key": "b5", "kind": "footer", "text": "Footer updated", "path": "", "hash": "h5", "simhash": 0, "weight": 0.2, "order": 4, "attrs": {}}]55    keys = {}56    for name, payload in (("t1", text1), ("t2", text2), ("t3", text3), ("b1", json.dumps(blocks1)), ("b2", json.dumps(blocks2)), ("b3", json.dumps(blocks3))):57        key, _size, _created = archive.put_text(payload)58        keys[name] = key59        d["objects"].append(key)6061    async with transaction() as conn:62        await execute(conn, "insert into countries (code, name, region, subregion, lat, lon) values ('ZZ', 'Ztestland', 'Test Region', 'Test Sub', 45.5, -73.6) "63                            "on conflict (code) do nothing")64        await execute(conn, "insert into industries (slug, name, description, keywords, sort_order) values (:s, 'Ztest Industry', 'Synthetic test industry', "65                            "array['ztest'], 999) on conflict (slug) do nothing", s=d["industry"])66        await execute(conn, "insert into connectors (id, name, version, category) values (:id, 'Ztest generic', '1', 'homepage') on conflict (id) do nothing", id=ids["connector"])67        for key, slug, name, imp, indexed in ((ids["alpha"], d["alpha_slug"], f"Ztest Alpha {SUFFIX}", 0.9, True), (ids["beta"], d["beta_slug"], f"Ztest Beta {SUFFIX}", 0.5, False)):68            await execute(conn, "insert into companies (id, slug, display_name, legal_name, canonical_domain, website, description, industries, industry_primary, country, "69                                "hq_city, public_company, importance, tier, indexed, status, onboarding_status, first_observed_at, last_observed_at, last_event_at) values "70                                "(:id, :slug, :name, :legal, :domain, :website, 'Synthetic company used by the API tests', cast(:inds as text[]), :ip, 'ZZ', 'Testville', "71                                "false, :imp, 2, :indexed, 'ACTIVE', 'active', :t0, :t1, :t1)",72                          id=key, slug=slug, name=name, legal=name + " Inc.", domain=f"{slug}.example", website=f"https://{slug}.example", inds=[d["industry"]],73                          ip=d["industry"], imp=imp, indexed=indexed, t0=now - timedelta(days=40), t1=now - timedelta(minutes=5))74        await execute(conn, "insert into company_aliases (company_id, alias, alias_norm, kind) values (:c, :a, :n, 'brand')", c=ids["alpha"], a=f"Alphaz {SUFFIX}", n=f"alphaz{SUFFIX}")75        await execute(conn, "insert into domains (id, company_id, domain, kind) values (:id, :c, :d, 'primary')", id=f"dom_ztestapi{SUFFIX}1", c=ids["alpha"], d=f"{d['alpha_slug']}.example")76        await execute(conn, "insert into company_relationships (id, from_company_id, to_company_id, kind, confidence) values (:id, :a, :b, 'PARTNER_OF', 0.8)",77                      id=f"rel_ztestapi{SUFFIX}1", a=ids["alpha"], b=ids["beta"])78        for sid, cid, surface, url, snaps, changes, events in ((ids["sensor_home"], ids["alpha"], "pricing", f"https://{d['alpha_slug']}.example/pricing", 3, 1, 2),79                                                                (ids["sensor_careers"], ids["alpha"], "careers", f"https://{d['alpha_slug']}.example/careers", 0, 0, 1),80                                                                (ids["sensor_beta"], ids["beta"], "homepage", f"https://{d['beta_slug']}.example/", 0, 0, 1)):81            await execute(conn, "insert into sensors (id, company_id, surface, connector_id, url, canonical_url, domain, status, tier, quality_score, base_interval_s, "82                                "current_interval_s, next_run_at, last_run_at, last_success_at, last_change_at, last_status, observation_count, snapshot_count, change_count, "83                                "meaningful_change_count, event_count, last_snapshot_id) values (:id, :cid, :surface, :conn, :url, :url, :domain, 'active', 'C', 80, 21600, 21600, "84                                ":next, :last, :last, :last, 200, :obs, :snaps, :changes, :changes, :events, :lastsnap)",85                          id=sid, cid=cid, surface=surface, conn=ids["connector"], url=url, domain=url.split("/")[2], next=now + timedelta(hours=1),86                          last=now - timedelta(minutes=10), obs=snaps + 4, snaps=snaps, changes=changes, events=events, lastsnap=ids["snap3"] if snaps else None)87        for sid_key, ver, tkey, bkey, prev, when in (("snap1", 1, "t1", "b1", None, now - timedelta(days=2)), ("snap2", 2, "t2", "b2", ids["snap1"], now - timedelta(days=1)),88                                                     ("snap3", 3, "t3", "b3", ids["snap2"], now - timedelta(hours=2))):89            await execute(conn, "insert into snapshots (id, sensor_id, company_id, previous_snapshot_id, version_no, fetched_at, content_hash, normalized_hash, structural_hash, "90                                "text_key, blocks_key, extracted, extracted_summary, title, language, text_length, block_count) values (:id, :sid, :cid, :prev, :ver, :when, :h, :h, :h, "91                                ":tk, :bk, cast(:ex as jsonb), cast(:sum as jsonb), 'Ztest Alpha — pricing', 'en', :tl, :bc)",92                          id=ids[sid_key], sid=ids["sensor_home"], cid=ids["alpha"], prev=prev, ver=ver, when=when, h=keys[tkey], tk=keys[tkey], bk=keys[bkey],93                          ex=jsonb({"plans": [{"plan_name": "Starter"}]}), sum=jsonb({"plan_count": ver + 1}), tl=len(text1), bc=3)94        diff = {"added": [{"key": "b4", "kind": "pricing_plan", "path": "Pricing > Team", "before": None, "after": "New: Team plan", "weight": 1.5, "similarity": None}],95                "removed": [], "modified": [{"key": "b2", "kind": "pricing_plan", "path": "Pricing > Starter", "before": "Starter plan $10 per month",96                                             "after": "Starter plan $12 per month", "weight": 1.5, "similarity": 0.9}], "moved": [],97                "counts": {"added": 1, "removed": 0, "modified": 1, "moved": 0}, "text_delta_ratio": 0.18, "similarity": 0.82, "reasons": ["pricing_plan modified"]}98        await execute(conn, "insert into changes (id, sensor_id, company_id, surface, snapshot_before, snapshot_after, detected_at, significance, kind, blocks_added, blocks_modified, "99                            "text_delta_ratio, similarity, diff, structured_delta, status) values (:id, :sid, :cid, 'pricing', :b, :a, :when, 0.72, 'major', 1, 1, 0.18, 0.82, "100                            "cast(:diff as jsonb), cast(:sd as jsonb), 'processed')",101                      id=ids["change"], sid=ids["sensor_home"], cid=ids["alpha"], b=ids["snap1"], a=ids["snap2"], when=now - timedelta(days=1), diff=jsonb(diff),102                      sd=jsonb({"plans": {"price_changed": [{"plan_name": "Starter", "before": 10, "after": 12, "currency": "USD", "billing_period": "month", "pct": 20}]}}))103        events = (104            (ids["ev_pricing"], ids["alpha"], ids["sensor_home"], ids["change"], "pricing", "PRICING", "PRICE_INCREASE", 0.8, 0.92, "HIGH_CONFIDENCE",105             "Starter plan price increased from $10 to $12 per month", "$10", "$12", now - timedelta(minutes=30), ["pricing"], "active"),106            (ids["ev_hiring"], ids["alpha"], ids["sensor_careers"], None, "careers", "HIRING", "JOB_COUNT_INCREASE", 0.6, 0.85, "HIGH_CONFIDENCE",107             "Open positions increased from 1 to 2", "1", "2", now - timedelta(hours=3), ["hiring", "ai"], "active"),108            (ids["ev_product"], ids["alpha"], ids["sensor_home"], None, "pricing", "PRODUCT", "NEW_PRODUCT", 0.7, 0.75, "LIKELY",109             "New plan listed: Team", None, "Team", now - timedelta(hours=6), ["product"], "active"),110            (ids["ev_leader"], ids["beta"], ids["sensor_beta"], None, "homepage", "LEADERSHIP", "NEW_EXECUTIVE", 0.8, 0.7, "LIKELY",111             "New executive listed: Jane Ztest (Chief Test Officer)", None, "Jane Ztest", now - timedelta(days=3), ["leadership"], "active"),112            (ids["ev_retracted"], ids["beta"], ids["sensor_beta"], None, "homepage", "OTHER", "OTHER", 0.2, 0.5, "INFERRED",113             "Retracted synthetic event", None, None, now - timedelta(days=4), [], "retracted"),114        )115        for eid, cid, sid, chg, surface, et, est, imp, conf, label, title, old, new, when, tags, status in events:116            await execute(conn, "insert into events (id, company_id, sensor_id, change_id, surface, event_type, event_subtype, importance, confidence, confidence_label, title, "117                                "summary, old_value, new_value, payload, entities, tags, detected_at, source_url, origin, status, dedupe_key) values (:id, :cid, :sid, :chg, :surface, "118                                ":et, :est, :imp, :conf, :label, :title, 'Detected on a monitored public page.', :old, :new, '{}'::jsonb, '{}'::jsonb, cast(:tags as text[]), :when, "119                                ":url, 'deterministic', :status, :dk)",120                          id=eid, cid=cid, sid=sid, chg=chg, surface=surface, et=et, est=est, imp=imp, conf=conf, label=label, title=title, old=old, new=new, tags=tags,121                          when=when, url=f"https://{d['alpha_slug']}.example/{surface}", status=status, dk=f"ztest:{eid}")122        await execute(conn, "insert into event_sources (event_id, sensor_id, source_url, snapshot_id, surface, kind) values (:e, :s, :u, :snap, 'pricing', 'primary')",123                      e=ids["ev_pricing"], s=ids["sensor_home"], u=f"https://{d['alpha_slug']}.example/pricing", snap=ids["snap2"])124        for i, (title, status, ai, country) in enumerate((("Senior ML Engineer", "open", True, "ZZ"), ("Account Executive", "open", False, "ZZ"),125                                                          ("Office Manager", "no_longer_listed", False, None))):126            await execute(conn, "insert into jobs (id, company_id, sensor_id, fingerprint, title, department, location_text, city, country, remote, employment_type, first_seen_at, "127                                "last_seen_at, removed_at, status, is_ai) values (:id, :cid, :sid, :fp, :title, 'Engineering', 'Testville', 'Testville', :country, :remote, 'full_time', "128                                ":seen, :last, :removed, :status, :ai)",129                          id=f"job_ztestapi{SUFFIX}{i}", cid=ids["alpha"], sid=ids["sensor_careers"], fp=f"ztest-{SUFFIX}-{i}", title=title, country=country, remote=(i == 0),130                          seen=now - timedelta(days=3 + i), last=now - timedelta(hours=1), removed=(now - timedelta(days=1)) if status != "open" else None, status=status, ai=ai)131        await execute(conn, "insert into people (id, company_id, name, name_norm, title, role_category, is_executive, status) values (:id, :cid, 'Jane Ztest', :n, 'Chief Test Officer', 'other', true, 'listed')",132                      id=f"person_ztestapi{SUFFIX}1", cid=ids["alpha"], n=f"janeztest{SUFFIX}")133        await execute(conn, "insert into people (id, company_id, name, name_norm, title, role_category, is_executive, status, removed_at) values (:id, :cid, 'John Former', :n, 'CFO', 'cfo', true, 'no_longer_listed', now())",134                      id=f"person_ztestapi{SUFFIX}2", cid=ids["alpha"], n=f"johnformer{SUFFIX}")135        await execute(conn, "insert into products (id, company_id, name, name_norm, category, status) values (:id, :cid, 'Ztest Widget', :n, 'widgets', 'listed')",136                      id=f"prod_ztestapi{SUFFIX}1", cid=ids["alpha"], n=f"ztestwidget{SUFFIX}")137        await execute(conn, "insert into pricing_plans (id, company_id, plan_name, plan_norm, currency, billing_period, price, price_text, features, status, version_no) values "138                            "(:id, :cid, 'Starter', 'starter', 'USD', 'month', 12, '$12 per month', '[\"1 seat\"]'::jsonb, 'current', 2)", id=f"plan_ztestapi{SUFFIX}2", cid=ids["alpha"])139        await execute(conn, "insert into pricing_plans (id, company_id, plan_name, plan_norm, currency, billing_period, price, price_text, status, version_no, valid_to) values "140                            "(:id, :cid, 'Starter', 'starter', 'USD', 'month', 10, '$10 per month', 'superseded', 1, now())", id=f"plan_ztestapi{SUFFIX}1", cid=ids["alpha"])141        await execute(conn, "insert into locations (id, company_id, kind, name, name_norm, city, country, lat, lon, status) values (:id, :cid, 'headquarters', 'Testville HQ', :n, 'Testville', 'ZZ', 45.5, -73.6, 'listed')",142                      id=f"loc_ztestapi{SUFFIX}1", cid=ids["alpha"], n=f"testvillehq{SUFFIX}")143        await execute(conn, "insert into news_items (id, company_id, url, canonical_url, title, category, published_at) values (:id, :cid, :u, :u, 'Ztest Alpha announces Team plan', 'press', now())",144                      id=f"news_ztestapi{SUFFIX}1", cid=ids["alpha"], u=f"https://{d['alpha_slug']}.example/news/team-plan")145        for cid, metric, value in ((ids["alpha"], "activity_score", 72.34), (ids["alpha"], "hiring_momentum_30d", 12.5), (ids["alpha"], "open_jobs", 2), (ids["alpha"], "ai_adoption", 40.0),146                                   (ids["alpha"], "product_velocity", 55.0), (ids["beta"], "activity_score", 30.0), (ids["beta"], "hiring_momentum_30d", -8.0)):147            await execute(conn, "insert into metrics_current (company_id, metric, value, confidence, inputs, formula_version) values (:cid, :m, :v, 0.8, '{\"events\": 3}'::jsonb, 'metrics-v1')",148                          cid=cid, m=metric, v=value)149        for i in range(10):150            await execute(conn, "insert into metric_series (company_id, metric, day, value, confidence, formula_version) values (:cid, 'activity_score', :day, :v, 0.8, 'metrics-v1')",151                          cid=ids["alpha"], day=(now - timedelta(days=9 - i)).date(), v=50 + i * 2)152        await execute(conn, "insert into company_daily (company_id, day, observations, changes, events) values (:cid, :day, 5, 1, 2)", cid=ids["alpha"], day=(now - timedelta(days=1)).date())153        await execute(conn, "insert into signals (id, company_id, scope, scope_key, kind, strength, confidence, title, explanation, window_days) values (:id, :cid, 'company', :slug, "154                            "'hiring_surge', 0.7, 0.6, 'Hiring surge signal', 'Open positions doubled over 7 days.', 7)", id=ids["signal"], cid=ids["alpha"], slug=d["alpha_slug"])155        await execute(conn, "insert into queue_jobs (id, kind, key, payload, status, attempts, max_attempts, last_error) values (:id, 'run_sensor', :key, '{}'::jsonb, 'dead', 3, 3, 'boom')",156                      id=ids["queue_dead"], key=f"ztest:{SUFFIX}:dead")157        await execute(conn, "insert into failures (id, sensor_id, company_id, failure_class, status_code, message, url) values (:id, :sid, :cid, 'HTTP_5XX', 503, 'synthetic', 'https://x.example')",158                      id=ids["failure"], sid=ids["sensor_beta"], cid=ids["beta"])159        await execute(conn, "insert into review_queue (id, kind, ref_id, company_id, payload) values (:id, 'major_event', :ref, :cid, '{}'::jsonb)",160                      id=ids["review"], ref=ids["ev_pricing"], cid=ids["alpha"])161        await execute(conn, "insert into api_keys (id, key_hash, prefix, name, tier) values (:id, :h, :p, 'ztest paid', 'paid')",162                      id=ids["api_key"], h=hashlib.sha256(RAW_API_KEY.encode()).hexdigest(), p=RAW_API_KEY[:12])163    return d164165166async def _cleanup(d: dict[str, Any]) -> None:167    async with transaction() as conn:168        await execute(conn, "delete from companies where slug like :p", p=f"ztest-api-%{d['suffix']}%")169        await execute(conn, "delete from companies where slug like 'ztest-api-created-%'")170        await execute(conn, "delete from companies where canonical_domain like 'ztest-created-%'")171        await execute(conn, "delete from queue_jobs where key like :p", p=f"ztest:{d['suffix']}%")172        await execute(conn, "delete from queue_jobs where key like 'discover:co_ztestapi%'")173        await execute(conn, "delete from api_keys where id = :id", id=d["api_key"])174        await execute(conn, "delete from owners where token_hash = :h", h=hashlib.sha256(OWNER_TOKEN.encode()).hexdigest())175        await execute(conn, "delete from connectors where id = :id", id=d["connector"])176        await execute(conn, "delete from industries where slug = :s", s=d["industry"])177        left = await fetch_val(conn, "select count(*) from companies where country = 'ZZ'")178        if not left:179            await execute(conn, "delete from countries where code = 'ZZ'")180    for key in d["objects"]:181        try:182            os.unlink(archive.object_path(key))183        except OSError:184            pass185    cache.clear()186187188_STATE: dict[str, Any] = {"refs": 0, "data": None}   # one dataset per process even though each test module re-exports the fixture189190191@pytest_asyncio.fixture(loop_scope="session", scope="session")192async def fixture_data():  # type: ignore[no-untyped-def]193    from companyatlas.api import ratelimit194195    if _STATE["data"] is None:196        cache.clear()197        _STATE["anon_limit"] = ratelimit.limiter.limits["anonymous"]198        ratelimit.limiter.limits["anonymous"] = 1_000_000       # the suite alone exceeds 120 req/min; the 429 path has its own test199        _STATE["data"] = await _insert_fixture()200    _STATE["refs"] += 1201    try:202        yield _STATE["data"]203    finally:204        _STATE["refs"] -= 1205        if _STATE["refs"] == 0:206            data, _STATE["data"] = _STATE["data"], None207            ratelimit.limiter.limits["anonymous"] = _STATE["anon_limit"]208            await _cleanup(data)209            await dispose()210211212@pytest_asyncio.fixture(loop_scope="session")213async def client(fixture_data):  # type: ignore[no-untyped-def]214    from companyatlas.api.main import app215216    async with httpx.AsyncClient(transport=httpx.ASGITransport(app=app), base_url="http://testserver", timeout=30) as c:217        yield c218219220__all__ = ["ADMIN", "OWNER", "OWNER_TOKEN", "RAW_API_KEY", "client", "fixture_data"]221