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 EMPTY_BATCH_ERROR, Pipeline, QualityError from datahub.serving import ApiError, V1API from datahub.settings import Settings from datahub.timeutil import SHANGHAI, isoformat from tests.fixtures import TRADE_DATE, fake_transport DATASET_API = { "daily": "daily", "valuation": "daily_basic", "moneyflow": "moneyflow", "auction": "stk_auction", "index_daily": "index_daily", } def empty_transport_for(*datasets: str): blocked = {DATASET_API[name] for name in datasets} if "daily" in datasets: blocked.add("adj_factor") def transport(api_name, params, fields): if api_name in blocked: return [] return fake_transport(api_name, params, fields) return 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) self.assertEqual(ref["calendar_from"], "20160101") 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) def test_empty_index_daily_keeps_previous_official_readable(self) -> None: pipe, db = make_pipeline() pipe.ingest_reference(TRADE_DATE) first = pipe.run_dataset("index_daily", TRADE_DATE) self.assertEqual(first["state"], "published") self.assertGreater(first["rows"], 0) pipe.adapter._transport = empty_transport_for("index_daily") with self.assertRaises(QualityError) as ctx: pipe.run_dataset("index_daily", TRADE_DATE) self.assertIn(EMPTY_BATCH_ERROR, ctx.exception.report["errors"]) pub = db.fetchone( "SELECT * FROM publications WHERE dataset='index_daily' AND trade_date=?", (TRADE_DATE,), ) self.assertEqual(pub["active_batch"], first["batch_id"]) self.assertEqual(pub["state"], "published") empty_batch = db.fetchone( "SELECT * FROM batches WHERE batch_id=?", (f"{TRADE_DATE}-index_daily-002",), ) self.assertIsNotNone(empty_batch) self.assertEqual(empty_batch["state"], "staged") self.assertEqual(empty_batch["rows_out"], 0) self.assertNotEqual(empty_batch["state"], "failed") self.assertIn(EMPTY_BATCH_ERROR, empty_batch["error"] or "") api = V1API(pipe.db, pipe, pipe.settings) payload = api.handle("/v1/indexes/bars", {"date": [TRADE_DATE], "code": ["000001.SH"]}) self.assertEqual(payload["meta"]["batch_id"], first["batch_id"]) self.assertEqual(payload["meta"]["tier"], "official") self.assertEqual(payload["meta"]["state"], "published") self.assertTrue(payload["data"]) def test_empty_index_daily_first_batch_is_not_official(self) -> None: empty_date = "20260902" pipe, db = make_pipeline() pipe.ingest_reference(TRADE_DATE) pipe.adapter._transport = empty_transport_for("index_daily") with self.assertRaises(QualityError) as ctx: pipe.run_dataset("index_daily", empty_date) self.assertIn(EMPTY_BATCH_ERROR, ctx.exception.report["errors"]) batch = db.fetchone("SELECT * FROM batches WHERE batch_id=?", (f"{empty_date}-index_daily-001",)) self.assertEqual(batch["state"], "staged") self.assertEqual(batch["rows_out"], 0) self.assertIsNone( db.fetchone( "SELECT * FROM publications WHERE dataset='index_daily' AND trade_date=?", (empty_date,), ) ) api = V1API(pipe.db, pipe, pipe.settings) with self.assertRaises(ApiError) as api_ctx: api.handle("/v1/indexes/bars", {"date": [empty_date], "code": ["000001.SH"]}) self.assertEqual(api_ctx.exception.code, "DATASET_NOT_PUBLISHED") def test_nonempty_index_daily_still_publishes(self) -> None: pipe, db = make_pipeline() pipe.ingest_reference(TRADE_DATE) result = pipe.run_dataset("index_daily", TRADE_DATE) self.assertEqual(result["state"], "published") self.assertEqual(result["rows"], 4) pub = db.fetchone( "SELECT * FROM publications WHERE dataset='index_daily' AND trade_date=?", (TRADE_DATE,), ) self.assertEqual(pub["active_batch"], result["batch_id"]) rows = db.fetchall("SELECT * FROM eod_index_bars WHERE batch_id=?", (result["batch_id"],)) self.assertEqual(len(rows), 4) def test_empty_batch_guard_covers_all_official_datasets(self) -> None: for dataset in ("daily", "valuation", "moneyflow", "auction", "index_daily"): with self.subTest(dataset=dataset): pipe, db = make_pipeline() pipe.ingest_reference(TRADE_DATE) first = pipe.run_dataset(dataset, TRADE_DATE) pipe.adapter._transport = empty_transport_for(dataset) with self.assertRaises(QualityError) as ctx: pipe.run_dataset(dataset, TRADE_DATE) self.assertTrue(ctx.exception.report["hard_fail"]) self.assertIn(EMPTY_BATCH_ERROR, ctx.exception.report["errors"]) pub = db.fetchone( "SELECT * FROM publications WHERE dataset=? AND trade_date=?", (dataset, TRADE_DATE), ) self.assertEqual(pub["active_batch"], first["batch_id"]) self.assertEqual(pub["state"], "published") empty_batch = db.fetchone( "SELECT * FROM batches WHERE dataset=? AND trade_date=? AND batch_id != ?", (dataset, TRADE_DATE, first["batch_id"]), ) self.assertEqual(empty_batch["state"], "staged") self.assertEqual(empty_batch["rows_out"], 0) def test_publish_entry_rejects_empty_staging_without_moving_pointer(self) -> None: pipe, db = make_pipeline() pipe.ingest_reference(TRADE_DATE) first = pipe.run_dataset("moneyflow", TRADE_DATE) batch_id = pipe.next_batch_id("moneyflow", TRADE_DATE) pipe._set_batch(batch_id, "moneyflow", TRADE_DATE, "publishing", 1, rows_in=0, rows_out=0) pipe._stage("moneyflow", batch_id, []) with self.assertRaises(QualityError) as ctx: pipe.publish("moneyflow", TRADE_DATE, batch_id, state="published") self.assertIn(EMPTY_BATCH_ERROR, ctx.exception.report["errors"]) pub = db.fetchone( "SELECT * FROM publications WHERE dataset='moneyflow' AND trade_date=?", (TRADE_DATE,), ) self.assertEqual(pub["active_batch"], first["batch_id"]) self.assertEqual(pub["state"], "published") history = db.fetchall( "SELECT batch_id FROM publication_history WHERE dataset='moneyflow' AND trade_date=?", (TRADE_DATE,), ) self.assertEqual({row["batch_id"] for row in history}, {first["batch_id"]}) if __name__ == "__main__": unittest.main()