428 lines
14 KiB
Python
428 lines
14 KiB
Python
"""stack market — entities, gaps."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import re
|
|
from datetime import date
|
|
from pathlib import Path
|
|
|
|
import duckdb
|
|
import httpx
|
|
import pytest
|
|
from typer.testing import CliRunner
|
|
|
|
import cli.market as mk
|
|
from bib.store import Store
|
|
from cli import app
|
|
from market import edgar, http
|
|
from market import tables as T
|
|
from market.seed import SEAT_LABEL_KEYS, load_seed
|
|
from pfs.codetables import SeatRow, write_seats
|
|
|
|
runner = CliRunner()
|
|
FX = Path(__file__).parent.parent / "market" / "fixtures"
|
|
|
|
|
|
class _Pacer:
|
|
def wait(self):
|
|
pass
|
|
|
|
|
|
def _seat(org, specialty, kind="specialty", year=2024):
|
|
return SeatRow("RUC", year, org, specialty, "permanent", "u", "K", "s", kind)
|
|
|
|
|
|
def _specialties() -> set[str]:
|
|
companies, orgs = load_seed()
|
|
out = {o.specialty for o in orgs if o.specialty}
|
|
out |= {
|
|
link.target
|
|
for e in (*companies, *orgs)
|
|
for link in e.links
|
|
if link.scope == "specialty"
|
|
}
|
|
return out - set(SEAT_LABEL_KEYS)
|
|
|
|
|
|
def _build_db() -> duckdb.DuckDBPyConnection:
|
|
c = duckdb.connect(":memory:")
|
|
c.execute("CREATE SCHEMA terminology")
|
|
c.execute(
|
|
"CREATE TABLE terminology.cms_acceptable_provider_specialty_codes "
|
|
"(specialty_description VARCHAR)"
|
|
)
|
|
c.executemany(
|
|
"INSERT INTO terminology.cms_acceptable_provider_specialty_codes VALUES (?)",
|
|
[(s.title(),) for s in sorted(_specialties())],
|
|
)
|
|
write_seats(
|
|
c,
|
|
"RUC",
|
|
[
|
|
_seat("Family Medicine", "FAMILY PRACTICE"),
|
|
_seat("Cardiology", "CARDIOLOGY"),
|
|
_seat("RUC Chair", "", "category"),
|
|
],
|
|
)
|
|
return c
|
|
|
|
|
|
class _Batch:
|
|
def __init__(self, c):
|
|
self.c = c
|
|
|
|
def __enter__(self):
|
|
return self.c
|
|
|
|
def __exit__(self, *a):
|
|
return False
|
|
|
|
|
|
@pytest.fixture
|
|
def env(monkeypatch, tmp_path):
|
|
c = _build_db()
|
|
monkeypatch.setattr(mk, "_batch", lambda: _Batch(c))
|
|
monkeypatch.setattr(mk, "_read", lambda: c.cursor())
|
|
store = Store(":memory:", storage_dir=tmp_path / "st")
|
|
monkeypatch.setattr(mk, "_store", lambda: store)
|
|
published: list[bool] = []
|
|
monkeypatch.setattr(mk, "_publish", lambda: published.append(True))
|
|
monkeypatch.setattr(http, "hoard_root", lambda: tmp_path / "market")
|
|
monkeypatch.setattr(edgar, "new_pacer", lambda: _Pacer())
|
|
try:
|
|
yield c, store, published, tmp_path / "market"
|
|
finally:
|
|
store.close()
|
|
c.close()
|
|
|
|
|
|
def _cache_offline(hoard: Path) -> None:
|
|
sec = hoard / "sec"
|
|
(sec / "submissions").mkdir(parents=True)
|
|
(sec / "company_tickers.json").write_bytes(
|
|
(FX / "company_tickers.json").read_bytes()
|
|
)
|
|
for name in (
|
|
"submissions_CIK0001564406.json",
|
|
"CIK0001564406-submissions-001.json",
|
|
):
|
|
target = name.replace("submissions_", "")
|
|
if name.startswith("submissions_"):
|
|
target = "CIK0001564406.json"
|
|
(sec / "submissions" / target).write_bytes((FX / name).read_bytes())
|
|
|
|
|
|
def _count(c, table):
|
|
return c.execute(f"SELECT count(*) FROM {table}").fetchone()[0]
|
|
|
|
|
|
def test_entities_offline_writes_and_logs(env):
|
|
c, store, published, hoard = env
|
|
_cache_offline(hoard)
|
|
# a stale curated link the seed no longer carries
|
|
T.write_links(
|
|
c, [T.LinkRow("0001564406", "program", "QPP-MIPS", "curated", "x", "url", "u")]
|
|
)
|
|
res = runner.invoke(app, ["market", "entities", "--offline", "--write"])
|
|
assert res.exit_code == 0, res.output
|
|
companies, orgs = load_seed()
|
|
assert _count(c, "market.entity") == len(companies) + len(orgs)
|
|
osh = next(e for e in T.read_entities(c, "company") if e.key == "0001564406")
|
|
assert osh.status == "acquired" and osh.ended_on == date(2023, 5, 2)
|
|
assert osh.sic == "8000" and osh.item_key
|
|
assert "no submissions for AMED" in res.output
|
|
# the five fixture filings, and the seed's dated ticker
|
|
assert len(T.read_filings(c, "0001564406")) == 5
|
|
assert [s.source for s in T.read_securities(c, "0001564406")] == ["seed"]
|
|
links = T.read_links(c, basis="curated")
|
|
assert all(l.item_key for l in links)
|
|
assert ("program", "QPP-MIPS") not in {(l.scope, l.target) for l in links}
|
|
assert [l.target for l in T.read_links(c, basis="sic")] == ["8000"]
|
|
logs = c.execute(
|
|
"SELECT table_name, rows FROM cms.ingest_log WHERE module = 'market.edgar' "
|
|
"ORDER BY table_name"
|
|
).fetchall()
|
|
assert [t for t, _ in logs] == [
|
|
"market.entity",
|
|
"market.entity_link",
|
|
"market.filing",
|
|
"market.security",
|
|
]
|
|
assert published == [True]
|
|
assert "retired 1" in res.output
|
|
assert store.count() > 0
|
|
|
|
|
|
def test_entities_rerun_is_idempotent(env):
|
|
c, _, published, hoard = env
|
|
_cache_offline(hoard)
|
|
for _ in range(2):
|
|
assert (
|
|
runner.invoke(app, ["market", "entities", "--offline", "--write"]).exit_code
|
|
== 0
|
|
)
|
|
companies, orgs = load_seed()
|
|
assert _count(c, "market.entity") == len(companies) + len(orgs)
|
|
assert _count(c, "market.filing") == 5
|
|
|
|
|
|
def test_entities_no_write_prints_counts(env, monkeypatch):
|
|
c, _, published, hoard = env
|
|
_cache_offline(hoard)
|
|
monkeypatch.setattr(
|
|
mk, "_batch", lambda: (_ for _ in ()).throw(AssertionError("no batch"))
|
|
)
|
|
res = runner.invoke(app, ["market", "entities", "--offline"])
|
|
assert res.exit_code == 0, res.output
|
|
companies, orgs = load_seed()
|
|
assert f"entities {len(companies) + len(orgs)}" in res.output
|
|
assert (
|
|
"company/acquired" in res.output
|
|
and "organization/specialty_society" in res.output
|
|
)
|
|
assert "dry run" in res.output and published == []
|
|
assert "warning" not in res.output # every seat specialty has a society
|
|
|
|
|
|
def test_entities_warns_about_uncovered_seat_specialty(env):
|
|
c, _, _, hoard = env
|
|
_cache_offline(hoard)
|
|
write_seats(c, "RUC", [_seat("Podiatry", "PODIATRY", year=2025)])
|
|
res = runner.invoke(app, ["market", "entities", "--offline"])
|
|
assert res.exit_code == 0
|
|
assert "seat specialties without a society: ['PODIATRY']" in res.output
|
|
|
|
|
|
def _handler(routes):
|
|
def handler(request: httpx.Request):
|
|
url = str(request.url)
|
|
for pattern, make in routes:
|
|
m = re.search(pattern, url)
|
|
if m:
|
|
return make(m)
|
|
return httpx.Response(404)
|
|
|
|
return handler
|
|
|
|
|
|
def _online(monkeypatch, routes):
|
|
monkeypatch.setattr(
|
|
http,
|
|
"client",
|
|
lambda source, key=None: httpx.Client(
|
|
transport=httpx.MockTransport(_handler(routes))
|
|
),
|
|
)
|
|
|
|
|
|
def test_entities_refusal_prints_gap(env, monkeypatch):
|
|
c, _, published, hoard = env
|
|
_online(monkeypatch, [("", lambda m: httpx.Response(403))])
|
|
res = runner.invoke(app, ["market", "entities", "--write"])
|
|
assert res.exit_code == 0, res.output
|
|
assert "gap: sec-edgar:entities:" in res.output and "403" in res.output
|
|
assert _count(c, "market.entity") == 0 and _count(c, "market.filing") == 0
|
|
assert [g.source for g in T.read_gaps(c)] == ["sec-edgar:entities"]
|
|
assert published == [True]
|
|
|
|
|
|
def test_entities_refusal_without_write_records_nothing(env, monkeypatch):
|
|
c, _, published, _ = env
|
|
_online(monkeypatch, [("", lambda m: httpx.Response(403))])
|
|
monkeypatch.setattr(
|
|
mk, "_batch", lambda: (_ for _ in ()).throw(AssertionError("no batch"))
|
|
)
|
|
res = runner.invoke(app, ["market", "entities"])
|
|
assert res.exit_code == 0 and "gap: sec-edgar:entities:" in res.output
|
|
assert published == []
|
|
|
|
|
|
def _sub_for(cik: str) -> dict:
|
|
return {
|
|
"cik": cik,
|
|
"name": f"Filer {cik}",
|
|
"sic": "8082",
|
|
"sicDescription": "Home Health",
|
|
"tickers": [],
|
|
"formerNames": [],
|
|
"filings": {
|
|
"recent": {
|
|
"accessionNumber": ["A-1", "A-2"],
|
|
"filingDate": ["2024-01-01", "2024-01-02"],
|
|
"reportDate": ["", ""],
|
|
"acceptanceDateTime": ["", ""],
|
|
"form": ["10-K", "4"],
|
|
"items": ["", ""],
|
|
"primaryDocument": ["a.htm", "b.xml"],
|
|
},
|
|
"files": [],
|
|
},
|
|
}
|
|
|
|
|
|
def test_sic_sweep_registers_filers_and_sic_links(env, monkeypatch):
|
|
c, _, published, _ = env
|
|
atom = (FX / "sic_8082.atom").read_bytes().replace(b"0001468328", b"0000000042")
|
|
seen: list[str] = []
|
|
|
|
def sub(m):
|
|
seen.append(m.group(1))
|
|
cik = m.group(1)
|
|
if cik == "0001564406":
|
|
return httpx.Response(
|
|
200, content=(FX / "submissions_CIK0001564406.json").read_bytes()
|
|
)
|
|
if cik == "0000000042":
|
|
return httpx.Response(404) # no submissions file: skipped
|
|
return httpx.Response(200, json=_sub_for(cik))
|
|
|
|
def page(m):
|
|
return httpx.Response(
|
|
200, content=(FX / "CIK0001564406-submissions-001.json").read_bytes()
|
|
)
|
|
|
|
_online(
|
|
monkeypatch,
|
|
[
|
|
(
|
|
r"company_tickers\.json",
|
|
lambda m: httpx.Response(
|
|
200, content=(FX / "company_tickers.json").read_bytes()
|
|
),
|
|
),
|
|
(r"SIC=\d+.*start=0", lambda m: httpx.Response(200, content=atom)),
|
|
(r"submissions-001\.json", page),
|
|
(r"CIK(\d{10})\.json", sub),
|
|
],
|
|
)
|
|
res = runner.invoke(app, ["market", "entities", "--sic-sweep", "--write"])
|
|
assert res.exit_code == 0, res.output
|
|
assert res.output.count("SIC ") == len(edgar.HEALTH_SICS)
|
|
keys = {e.key for e in T.read_entities(c, "company")}
|
|
assert {"0001725255", "0001468328"} <= keys and "0000000042" not in keys
|
|
# a sweep filer keeps only the sweep forms; the seed's filings keep everything
|
|
assert [f.form for f in T.read_filings(c, "0001725255")] == ["10-K"]
|
|
assert len(T.read_filings(c, "0001468328")) == 2
|
|
assert len(T.read_filings(c, "0001564406")) == 5
|
|
sic = T.read_links(c, basis="sic")
|
|
assert {l.target for l in sic} == {"8082", "8000"}
|
|
assert any(l.entity_key == "0001725255" for l in sic)
|
|
|
|
|
|
def test_entities_seed_errors_abort(env, monkeypatch):
|
|
c, _, published, _ = env
|
|
companies, orgs = load_seed()
|
|
bad = companies[0].__class__(
|
|
**{
|
|
**companies[0].__dict__,
|
|
"links": (
|
|
companies[0].links[0].__class__("family", "NOPE", companies[0].cite),
|
|
),
|
|
}
|
|
)
|
|
monkeypatch.setattr(mk, "load_seed", lambda: ((bad, *companies[1:]), orgs))
|
|
res = runner.invoke(app, ["market", "entities", "--offline", "--write"])
|
|
assert res.exit_code == 1 and "seed error: OSH" in res.output
|
|
assert c.execute(
|
|
"SELECT count(*) FROM information_schema.tables WHERE table_schema = 'market'"
|
|
).fetchone() == (0,)
|
|
assert published == []
|
|
|
|
|
|
def test_entities_missing_vocabulary_tables_abort(env, monkeypatch):
|
|
bare = duckdb.connect(":memory:")
|
|
monkeypatch.setattr(mk, "_read", lambda: bare)
|
|
res = runner.invoke(app, ["market", "entities", "--offline"])
|
|
assert res.exit_code == 1 and "cannot validate the seed" in res.output
|
|
bare.close()
|
|
|
|
|
|
def test_entities_other_read_errors_propagate(env, monkeypatch):
|
|
def boom():
|
|
raise RuntimeError("disk on fire")
|
|
|
|
monkeypatch.setattr(mk, "_read", boom)
|
|
res = runner.invoke(app, ["market", "entities", "--offline"])
|
|
assert isinstance(res.exception, RuntimeError)
|
|
|
|
|
|
def test_require_key_missing_writes_only_gap(env, monkeypatch):
|
|
c, _, published, _ = env
|
|
monkeypatch.delenv("TIINGO_API_KEY", raising=False)
|
|
assert mk._require_key("tiingo", "tiingo:prices", write=True) is None
|
|
assert [(g.source, g.note) for g in T.read_gaps(c)] == [
|
|
("tiingo:prices", "TIINGO_API_KEY not in .env")
|
|
]
|
|
assert all(_count(c, f"market.{t}") == 0 for t in ("entity", "security", "filing"))
|
|
assert published == [True]
|
|
|
|
|
|
def test_require_key_present_returns_it(env, monkeypatch):
|
|
monkeypatch.setenv("TIINGO_API_KEY", "abc")
|
|
assert mk._require_key("tiingo", "tiingo:prices", write=True) == "abc"
|
|
|
|
|
|
def test_require_key_missing_prints_gap_without_write(env, monkeypatch, capsys):
|
|
monkeypatch.delenv("HUD_API_TOKEN", raising=False)
|
|
assert mk._require_key("hud", "hud:hud", write=False) is None
|
|
assert "gap: hud:hud: HUD_API_TOKEN not in .env" in capsys.readouterr().out
|
|
c, _, published, _ = env
|
|
assert published == []
|
|
|
|
|
|
def test_gap_cleared_after_success(env):
|
|
c, _, _, hoard = env
|
|
_cache_offline(hoard)
|
|
T.write_gap(c, "sec-edgar:entities", "403")
|
|
T.write_gap(c, "tiingo:prices", "no key")
|
|
assert (
|
|
runner.invoke(app, ["market", "entities", "--offline", "--write"]).exit_code
|
|
== 0
|
|
)
|
|
assert [g.source for g in T.read_gaps(c)] == ["tiingo:prices"]
|
|
|
|
|
|
def test_gaps_command_lists_rows(env):
|
|
c, _, _, _ = env
|
|
T.write_gap(c, "tiingo:prices", "TIINGO_API_KEY not in .env")
|
|
res = runner.invoke(app, ["market", "gaps"])
|
|
assert res.exit_code == 0
|
|
assert "gap: tiingo:prices: TIINGO_API_KEY not in .env" in res.output
|
|
|
|
|
|
def test_gaps_command_empty_and_missing_tables(env, monkeypatch):
|
|
c, _, _, _ = env
|
|
T.ensure_tables(c)
|
|
assert "no gaps recorded" in runner.invoke(app, ["market", "gaps"]).output
|
|
bare = duckdb.connect(":memory:")
|
|
monkeypatch.setattr(mk, "_read", lambda: bare)
|
|
assert "no market tables yet" in runner.invoke(app, ["market", "gaps"]).output
|
|
|
|
class _Broken:
|
|
def execute(self, *a, **k):
|
|
raise RuntimeError("x")
|
|
|
|
def close(self):
|
|
pass
|
|
|
|
monkeypatch.setattr(mk, "_read", lambda: _Broken())
|
|
assert isinstance(runner.invoke(app, ["market", "gaps"]).exception, RuntimeError)
|
|
bare.close()
|
|
|
|
|
|
def test_default_indirections_resolve_lazily(monkeypatch):
|
|
import conf.connect as cc
|
|
|
|
calls: list[str] = []
|
|
monkeypatch.setattr(cc, "duckdb_batch", lambda n: calls.append(f"batch {n}") or n)
|
|
monkeypatch.setattr(
|
|
cc, "duckdb", lambda n, read_only=True: calls.append(f"read {n}") or n
|
|
)
|
|
monkeypatch.setattr(cc, "bib", lambda: calls.append("bib") or "store")
|
|
monkeypatch.setattr(
|
|
cc, "publish_replica", lambda n: calls.append(f"pub {n}") or "p"
|
|
)
|
|
assert mk._batch() == "aco" and mk._read() == "aco" and mk._store() == "store"
|
|
mk._publish()
|
|
assert calls == ["batch aco", "read aco", "bib", "pub aco"]
|