a836cda1b2
Co-authored-by: Cursor <cursoragent@cursor.com> Co-authored-by: multica-agent <github@multica.ai>
291 lines
12 KiB
Python
291 lines
12 KiB
Python
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()
|