Files
xiaobai-review/xiaobai-datahub/tests/test_pipeline.py
T
施工员andmultica-agent 3eaa36a8d5 feat(HEL-560): 数据中枢接管数据源/模型池/会员,注册改一次性邀请码
主站
- 新增 m0006 invite_codes 迁移;注册强制邀请码(首个管理员除外),消码与建号
  同一事务,并发提交只有一个能成功
- 新增 /api/hub-admin/* 服务端点(共享 HUB_ADMIN_TOKEN,先于鉴权校验),供数据
  中枢桥接读写会话/密码/模型池/会员/邀请码,并提供供应商模型列表拉取
- 前端:注册表单加邀请码(桌面 login、index.html、移动端);「系统管理」改为
  「数据中枢」入口指向 8766,原模型池与会员管理分区移除,仅留「行情管理」;
  随之清理陈旧 CSS

数据中枢
- 取消独立账号:删除 hub_admin/hub_sessions 与登录、改密、锁定逻辑,改为校验
  主站 xiaobai_session,仅管理员可进,CSRF 由会话派生,危险操作二次确认走主站
- 控制台新增数据源凭证可编辑区(原有内容一项不删)、供应商制模型池(自动拉取
  /models,失败退回卡内手动录入)、会员管理与邀请码页
- 日夜双主题:颜色收敛为同名 token 换值,SVG 改用 inline style 以吃到变量

自测
- 主站 verify_baseline 通过(498 项);数据中枢 235 项通过
- tools/verify_datahub_console.py 端到端跑通两服务真实对话;
  tools/verify_datahub_console_ui.py 浏览器跑通门禁/凭证/模型池/会员/主题/1030 窄屏

Co-authored-by: multica-agent <github@multica.ai>
2026-09-16 11:44:09 +08:00

290 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,
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()