259 lines
9.3 KiB
Python
259 lines
9.3 KiB
Python
"""Orchestration tests for the source-agnostic import framework.
|
|
|
|
These drive ``run_import`` with a fake importer to prove the framework's
|
|
guarantees: idempotent no_op on an unchanged revision, atomic promotion on a
|
|
new revision, and — the load-bearing one — a failed validation or a mid-run
|
|
exception leaves the live tables untouched.
|
|
|
|
The advisory-lock branch is a no-op on SQLite, so lock mutual-exclusion has NO
|
|
coverage here (PG-verify-pending); only the deterministic key derivation is
|
|
unit-tested. Per the SQLite StaticPool caveat we never share a connection: each
|
|
test uses its own temp-file engine and seeds/asserts with short-lived sessions
|
|
sequenced around the ``run_import`` call.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import os
|
|
import tempfile
|
|
from datetime import date, datetime, timezone
|
|
|
|
import pytest
|
|
from sqlalchemy import func, select
|
|
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine
|
|
|
|
from app.database import Base
|
|
import app.models # noqa: F401 register models on Base.metadata
|
|
from app.models.data_import_run import DataImportRun
|
|
from app.models.fundamental_snapshot import FundamentalSnapshot
|
|
from app.models.system_event import SystemEvent
|
|
from app.services.data_import import (
|
|
STATUS_DEFERRED,
|
|
STATUS_FAILED,
|
|
STATUS_NO_OP,
|
|
STATUS_PROMOTED,
|
|
ValidationResult,
|
|
_advisory_key,
|
|
run_import,
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
async def engine():
|
|
"""A dedicated temp-file SQLite engine (independent connections, unlike the
|
|
shared in-memory test engine) so ``run_import`` can pin its own connection."""
|
|
fd, path = tempfile.mkstemp(suffix=".db")
|
|
os.close(fd)
|
|
eng = create_async_engine(f"sqlite+aiosqlite:///{path}")
|
|
async with eng.begin() as conn:
|
|
await conn.run_sync(Base.metadata.create_all)
|
|
try:
|
|
yield eng
|
|
finally:
|
|
await eng.dispose()
|
|
try:
|
|
os.unlink(path)
|
|
except OSError:
|
|
pass
|
|
|
|
|
|
def _factory(engine) -> async_sessionmaker[AsyncSession]:
|
|
return async_sessionmaker(engine, class_=AsyncSession, expire_on_commit=False)
|
|
|
|
|
|
class FakeImporter:
|
|
"""Writes ``n_rows`` fundamental_snapshots on promote (CIK-keyed, no ticker
|
|
FK) so the live-table effect is easy to count."""
|
|
|
|
source = "sec_facts"
|
|
|
|
def __init__(self, revision, *, ok=True, retryable=False, n_rows=3, raise_in="none"):
|
|
self.revision = revision
|
|
self.ok = ok
|
|
self.retryable = retryable
|
|
self.n_rows = n_rows
|
|
self.raise_in = raise_in
|
|
self.staged_called = False
|
|
self.promoted = False
|
|
|
|
async def detect_revision(self, db):
|
|
if self.raise_in == "detect":
|
|
raise RuntimeError("boom-detect")
|
|
return self.revision
|
|
|
|
async def stage(self, db):
|
|
self.staged_called = True
|
|
if self.raise_in == "stage":
|
|
raise RuntimeError("boom-stage")
|
|
if self.raise_in == "cancel":
|
|
raise asyncio.CancelledError()
|
|
return list(range(self.n_rows))
|
|
|
|
async def validate(self, db, staged):
|
|
return ValidationResult(
|
|
ok=self.ok,
|
|
summary={"staged_rows": len(staged)},
|
|
source_max_date=date(2026, 7, 21),
|
|
messages=[] if self.ok else ["coverage below threshold"],
|
|
retryable=self.retryable,
|
|
)
|
|
|
|
async def promote(self, db, staged, run_id):
|
|
if self.raise_in == "promote":
|
|
# write one row THEN raise, to prove rollback undoes partial writes
|
|
db.add(_snapshot(self.revision, 999))
|
|
raise RuntimeError("boom-promote")
|
|
for i in staged:
|
|
db.add(_snapshot(self.revision, i, run_id=run_id))
|
|
self.promoted = True
|
|
self.promoted_run_id = run_id
|
|
return {"fundamental_snapshots": len(staged)}
|
|
|
|
|
|
def _snapshot(revision: str, i: int, run_id: int | None = None) -> FundamentalSnapshot:
|
|
return FundamentalSnapshot(
|
|
cik=f"{i:010d}",
|
|
accession=f"{revision}-{i:06d}",
|
|
form="10-Q",
|
|
filed_date=date(2026, 7, 1),
|
|
accepted_at=datetime(2026, 7, 1, tzinfo=timezone.utc),
|
|
period_end=date(2026, 6, 30),
|
|
fiscal_year=2026,
|
|
fiscal_period="Q2",
|
|
revenue=1000.0 + i,
|
|
import_run_id=run_id,
|
|
)
|
|
|
|
|
|
async def _count(factory, model) -> int:
|
|
async with factory() as s:
|
|
return (await s.execute(select(func.count()).select_from(model))).scalar_one()
|
|
|
|
|
|
async def _runs(factory) -> list[DataImportRun]:
|
|
async with factory() as s:
|
|
return list(
|
|
(await s.execute(select(DataImportRun).order_by(DataImportRun.id))).scalars()
|
|
)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_advisory_key_deterministic_and_distinct():
|
|
assert _advisory_key("sec_facts") == _advisory_key("sec_facts")
|
|
assert _advisory_key("sec_facts") != _advisory_key("dolt_earnings")
|
|
for src in ("sec_facts", "dolt_earnings", "dolt_stocks"):
|
|
k = _advisory_key(src)
|
|
assert -(2**63) <= k < 2**63 # fits Postgres bigint
|
|
|
|
|
|
async def test_promote_writes_and_records_run(engine):
|
|
factory = _factory(engine)
|
|
run = await run_import(FakeImporter("rev1", n_rows=4), engine=engine)
|
|
|
|
assert run is not None and run.status == STATUS_PROMOTED
|
|
assert run.revision == "rev1"
|
|
assert run.row_counts_json is not None and "fundamental_snapshots" in run.row_counts_json
|
|
assert run.source_max_date == date(2026, 7, 21)
|
|
assert run.completed_at is not None
|
|
assert await _count(factory, FundamentalSnapshot) == 4
|
|
# rows stamped with the run id
|
|
async with factory() as s:
|
|
stamped = (
|
|
await s.execute(select(FundamentalSnapshot.import_run_id))
|
|
).scalars().all()
|
|
assert stamped and all(rid == run.id for rid in stamped)
|
|
|
|
|
|
async def test_no_op_on_repeated_revision(engine):
|
|
factory = _factory(engine)
|
|
await run_import(FakeImporter("rev1", n_rows=4), engine=engine)
|
|
|
|
second = FakeImporter("rev1", n_rows=4)
|
|
run = await run_import(second, engine=engine)
|
|
|
|
assert run is not None and run.status == STATUS_NO_OP
|
|
assert second.staged_called is False # never fetched
|
|
assert await _count(factory, FundamentalSnapshot) == 4 # unchanged
|
|
|
|
runs = await _runs(factory)
|
|
assert [r.status for r in runs] == [STATUS_PROMOTED, STATUS_NO_OP]
|
|
|
|
|
|
async def test_new_revision_after_promote_stages_again(engine):
|
|
factory = _factory(engine)
|
|
await run_import(FakeImporter("rev1", n_rows=2), engine=engine)
|
|
run = await run_import(FakeImporter("rev2", n_rows=3), engine=engine)
|
|
|
|
assert run is not None and run.status == STATUS_PROMOTED
|
|
assert await _count(factory, FundamentalSnapshot) == 5 # 2 + 3
|
|
|
|
|
|
async def test_failed_validation_leaves_data_untouched(engine):
|
|
factory = _factory(engine)
|
|
await run_import(FakeImporter("rev1", n_rows=3), engine=engine) # baseline
|
|
|
|
run = await run_import(FakeImporter("rev2", ok=False, n_rows=5), engine=engine)
|
|
|
|
assert run is not None and run.status == STATUS_FAILED
|
|
assert "coverage" in (run.error_details or "")
|
|
assert await _count(factory, FundamentalSnapshot) == 3 # untouched
|
|
assert await _count(factory, SystemEvent) == 1 # alerted
|
|
|
|
|
|
async def test_retryable_validation_defers_without_alerting(engine):
|
|
factory = _factory(engine)
|
|
await run_import(FakeImporter("rev1", n_rows=3), engine=engine)
|
|
|
|
run = await run_import(
|
|
FakeImporter("rev2", ok=False, retryable=True, n_rows=5), engine=engine
|
|
)
|
|
|
|
assert run is not None and run.status == STATUS_DEFERRED
|
|
assert "coverage" in (run.error_details or "")
|
|
assert await _count(factory, FundamentalSnapshot) == 3 # untouched
|
|
assert await _count(factory, SystemEvent) == 0 # expected retry does not alert
|
|
|
|
|
|
async def test_exception_in_promote_rolls_back(engine):
|
|
factory = _factory(engine)
|
|
await run_import(FakeImporter("rev1", n_rows=3), engine=engine) # baseline
|
|
|
|
run = await run_import(FakeImporter("rev2", raise_in="promote"), engine=engine)
|
|
|
|
assert run is not None and run.status == STATUS_FAILED
|
|
assert "boom-promote" in (run.error_details or "")
|
|
assert await _count(factory, FundamentalSnapshot) == 3 # partial write rolled back
|
|
assert await _count(factory, SystemEvent) == 1
|
|
|
|
|
|
async def test_detection_failure_records_and_alerts(engine):
|
|
"""The external revision probe is the most likely failure — it must produce a
|
|
recorded failed run + alert, not an unrecorded escaping exception."""
|
|
factory = _factory(engine)
|
|
imp = FakeImporter("rev1", raise_in="detect")
|
|
|
|
run = await run_import(imp, engine=engine)
|
|
|
|
assert run is not None and run.status == STATUS_FAILED
|
|
assert "boom-detect" in (run.error_details or "")
|
|
assert imp.staged_called is False
|
|
assert await _count(factory, FundamentalSnapshot) == 0
|
|
assert await _count(factory, SystemEvent) == 1 # alerted
|
|
runs = await _runs(factory)
|
|
assert len(runs) == 1 and runs[0].status == STATUS_FAILED # attempt recorded
|
|
|
|
|
|
async def test_cancellation_marks_failed_and_reraises(engine):
|
|
factory = _factory(engine)
|
|
|
|
with pytest.raises(asyncio.CancelledError):
|
|
await run_import(FakeImporter("rev1", raise_in="cancel"), engine=engine)
|
|
|
|
runs = await _runs(factory)
|
|
assert len(runs) == 1 and runs[0].status == STATUS_FAILED # no lingering running row
|
|
assert runs[0].error_details == "cancelled"
|
|
assert await _count(factory, FundamentalSnapshot) == 0
|