xiaobai-review/xiaobai-datahub/tests/test_pipeline.py

from __future__ import annotations

import tempfile
import unittest
from datetime import datetime, timedelta
from pathlib import Path

from datahub.adapters.tushare import TushareAdapter
from datahub.crypto import SecretVault
from datahub.db import HubDB
from datahub.pipeline import Pipeline, QualityError
from datahub.settings import Settings
from datahub.timeutil import SHANGHAI, isoformat
from tests.fixtures import TRADE_DATE, fake_transport


def make_pipeline(before_commit=None, clock=None, quality=None) -> tuple[Pipeline, HubDB]:
    tmp = tempfile.TemporaryDirectory()
    db = HubDB(Path(tmp.name) / "hub.db")
    adapter = TushareAdapter("test-token", transport=fake_transport)
    quality_cfg = {
        "daily_row_ratio": 0.98,
        "null_rate_max": 0.01,
        "max_publish_attempts": 3,
        "publication_generations": 3,
        "job_run_retain_days": 90,
        "staging_retain_days": 14,
    }
    if quality:
        quality_cfg.update(quality)
    settings = Settings(
        encryption_key=SecretVault.generate_key(),
        api_token="t" * 32,
        admin_password="admin-pass",
        tushare_token="test-token",
        db_path=db.path,
        quality=quality_cfg,
        scheduler_enabled=False,
    )
    pipe = Pipeline(db, adapter, settings, before_commit=before_commit, clock=clock)
    pipe._tmp = tmp  # keep alive
    return pipe, db


class PipelineTests(unittest.TestCase):
    def test_reference_and_daily_publish(self) -> None:
        pipe, db = make_pipeline()
        ref = pipe.ingest_reference(TRADE_DATE)
        self.assertEqual(ref["stocks"], 2)
        result = pipe.run_dataset("daily", TRADE_DATE)
        self.assertEqual(result["state"], "published")
        self.assertEqual(result["rows"], 2)
        pub = db.fetchone("SELECT * FROM publications WHERE dataset='daily' AND trade_date=?", (TRADE_DATE,))
        self.assertEqual(pub["active_batch"], result["batch_id"])
        rows = db.fetchall("SELECT * FROM eod_bars WHERE batch_id=?", (result["batch_id"],))
        self.assertEqual(len(rows), 2)
        self.assertEqual(rows[0]["amount"] if rows[0]["ts_code"] == "600000.SH" else rows[1]["amount"], 2_000_000.0)

    def test_atomic_publish_abort_leaves_no_half_batch(self) -> None:
        pipe, db = make_pipeline()
        pipe.ingest_reference(TRADE_DATE)
        first = pipe.run_dataset("daily", TRADE_DATE)
        boom = {"n": 0}

        def explode() -> None:
            boom["n"] += 1
            raise RuntimeError("killed")

        pipe.before_commit = explode
        with self.assertRaises(RuntimeError):
            pipe.run_dataset("daily", TRADE_DATE)
        pub = db.fetchone("SELECT * FROM publications WHERE dataset='daily' AND trade_date=?", (TRADE_DATE,))
        self.assertEqual(pub["active_batch"], first["batch_id"])
        visible = db.fetchall(
            "SELECT DISTINCT batch_id FROM eod_bars WHERE trade_date=? AND batch_id=?",
            (TRADE_DATE, pub["active_batch"]),
        )
        self.assertEqual(len(visible), 1)

    def test_rollback_switches_active_batch(self) -> None:
        pipe, _db = make_pipeline()
        pipe.ingest_reference(TRADE_DATE)
        first = pipe.run_dataset("daily", TRADE_DATE)
        second = pipe.run_dataset("daily", TRADE_DATE)
        self.assertNotEqual(first["batch_id"], second["batch_id"])
        rolled = pipe.rollback("daily", TRADE_DATE, actor="test")
        self.assertEqual(rolled["active_batch"], first["batch_id"])
        from datahub.serving import V1API

        api = V1API(pipe.db, pipe, pipe.settings)
        payload = api.handle("/v1/bars/daily", {"date": [TRADE_DATE], "code": ["600000.SH"]})
        self.assertEqual(payload["meta"]["batch_id"], first["batch_id"])

    def test_row_ratio_gate_rejects_short_batch(self) -> None:
        pipe, _db = make_pipeline()
        pipe.ingest_reference(TRADE_DATE)
        original = fake_transport

        def short(api_name, params, fields):
            rows = original(api_name, params, fields)
            if api_name == "daily":
                return rows[:1]
            return rows

        pipe.adapter._transport = short
        with self.assertRaises(QualityError) as ctx:
            pipe.run_dataset("daily", TRADE_DATE)
        self.assertTrue(ctx.exception.report["hard_fail"])
        pub = pipe.db.fetchone("SELECT * FROM publications WHERE dataset='daily'")
        self.assertIsNone(pub)

    def test_wal_mode(self) -> None:
        pipe, db = make_pipeline()
        with db.connect() as connection:
            mode = connection.execute("PRAGMA journal_mode").fetchone()[0]
        self.assertEqual(str(mode).lower(), "wal")

    def test_cleanup_iso_timestamps_respect_retention_on_job_and_src(self) -> None:
        # job_runs.started_at / src_calls.created_at 存 ISO；旧实现用 YYYYMMDD 比较会误删同年记录。
        frozen = datetime(2026, 9, 2, 0, 30, tzinfo=SHANGHAI)
        retain_days = 90
        pipe, db = make_pipeline(clock=lambda: frozen, quality={"job_run_retain_days": retain_days})
        samples = {
            "today": isoformat(frozen),
            "within": isoformat(frozen - timedelta(days=retain_days - 1)),
            "expired": isoformat(frozen - timedelta(days=retain_days + 1)),
        }
        with db.write() as connection:
            for job_id, stamp in samples.items():
                connection.execute(
                    "INSERT INTO job_runs(job_id, state, started_at, attempt) VALUES (?,?,?,1)",
                    (job_id, "ok", stamp),
                )
                connection.execute(
                    "INSERT INTO src_calls(provider, endpoint, ok, latency_ms, error, created_at) VALUES (?,?,?,?,?,?)",
                    ("tushare", job_id, 1, 10, None, stamp),
                )

        pipe.cleanup()

        jobs = {row["job_id"] for row in db.fetchall("SELECT job_id FROM job_runs")}
        calls = {row["endpoint"] for row in db.fetchall("SELECT endpoint FROM src_calls")}
        kept = {"today", "within"}
        self.assertEqual(jobs, kept)
        self.assertEqual(calls, kept)


if __name__ == "__main__":
    unittest.main()