signal-platform/tests/unit/test_data_import_framework.py

"""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_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, n_rows=3, raise_in="none"):
        self.revision = revision
        self.ok = ok
        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"],
        )

    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_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