75c2e33b68
eod-refresh --force 与管理后台补数不再单数据集发布， 统一走 force_republish_boundary，避免绕过 A/B 完整边界。 Co-authored-by: Cursor <cursoragent@cursor.com> Co-authored-by: multica-agent <github@multica.ai>
373 lines
18 KiB
Python
373 lines
18 KiB
Python
from __future__ import annotations
|
|
|
|
import unittest
|
|
from pathlib import Path
|
|
import tempfile
|
|
|
|
from datahub.adapters.tushare import TushareAdapter
|
|
from datahub.crypto import SecretVault
|
|
from datahub.db import HubDB
|
|
from datahub.pipeline import Pipeline
|
|
from datahub.settings import Settings
|
|
from datahub.serving import V1API
|
|
from tests.fixtures import TRADE_DATE, fake_transport
|
|
|
|
GROUP_A = ("daily", "valuation", "moneyflow", "auction")
|
|
|
|
|
|
class GroupTransport:
|
|
"""fake_transport with per-API degradation switches for release-group tests."""
|
|
|
|
def __init__(self) -> None:
|
|
self.empty: set[str] = set()
|
|
self.keep_rows: dict[str, int] = {}
|
|
self.stocks: list[dict] | None = None
|
|
self.calls: list[str] = []
|
|
|
|
def __call__(self, api_name: str, params: dict, fields: str):
|
|
self.calls.append(api_name)
|
|
if api_name in self.empty:
|
|
return []
|
|
if api_name == "stock_basic" and self.stocks is not None:
|
|
return [dict(row) for row in self.stocks]
|
|
rows = fake_transport(api_name, params, fields)
|
|
keep = self.keep_rows.get(api_name)
|
|
if keep is not None:
|
|
return rows[:keep]
|
|
return rows
|
|
|
|
|
|
def make_pipe(transport: GroupTransport, quality_extra: dict | None = None):
|
|
tmp = tempfile.TemporaryDirectory()
|
|
db = HubDB(Path(tmp.name) / "hub.db")
|
|
adapter = TushareAdapter("test-token", transport=transport)
|
|
quality = {
|
|
"daily_row_ratio": 0.98,
|
|
"null_rate_max": 0.01,
|
|
"max_publish_attempts": 2,
|
|
"publication_generations": 3,
|
|
}
|
|
if quality_extra:
|
|
quality.update(quality_extra)
|
|
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,
|
|
scheduler_enabled=False,
|
|
)
|
|
pipe = Pipeline(db, adapter, settings)
|
|
pipe._tmp = tmp
|
|
return pipe, db
|
|
|
|
|
|
def publications_map(db: HubDB, day: str) -> dict[str, str]:
|
|
rows = db.fetchall("SELECT dataset, active_batch FROM publications WHERE trade_date = ?", (day,))
|
|
return {str(row["dataset"]): str(row["active_batch"]) for row in rows}
|
|
|
|
|
|
class ReleaseGroupSwitchTests(unittest.TestCase):
|
|
def setUp(self) -> None:
|
|
self.transport = GroupTransport()
|
|
self.pipe, self.db = make_pipe(self.transport)
|
|
self.pipe.ingest_reference(TRADE_DATE)
|
|
|
|
def test_whole_group_switches_in_one_publish_instant(self) -> None:
|
|
results = self.pipe.run_eod_batch_a(TRADE_DATE)
|
|
self.assertEqual(set(results), {*GROUP_A, "stocks"})
|
|
self.assertEqual({item["state"] for item in results.values()}, {"published"})
|
|
pubs = self.db.fetchall("SELECT * FROM publications WHERE trade_date = ?", (TRADE_DATE,))
|
|
self.assertEqual(len(pubs), 5)
|
|
self.assertEqual(len({row["published_at"] for row in pubs}), 1)
|
|
# official rows copied and serving resolves the new batches
|
|
api = V1API(self.db, self.pipe, self.pipe.settings)
|
|
payload = api.handle("/v1/bars/daily", {"date": [TRADE_DATE], "code": ["600000.SH"]})
|
|
self.assertEqual(payload["meta"]["batch_id"], results["daily"]["batch_id"])
|
|
stocks = api.handle("/v1/stocks", {})
|
|
self.assertEqual(stocks["meta"]["batch_id"], results["stocks"]["batch_id"])
|
|
|
|
def test_any_member_failure_blocks_entire_group(self) -> None:
|
|
self.transport.empty = {"daily_basic"} # valuation upstream returns nothing
|
|
results = self.pipe.run_eod_batch_a(TRADE_DATE)
|
|
self.assertEqual(results["valuation"]["state"], "failed")
|
|
self.assertEqual(results["moneyflow"]["state"], "aborted")
|
|
self.assertEqual(results["auction"]["state"], "aborted")
|
|
self.assertEqual(results["daily"]["state"], "failed") # staged fine, then abandoned
|
|
# nothing became visible, and the reason is recorded
|
|
self.assertEqual(publications_map(self.db, TRADE_DATE), {})
|
|
abandoned = self.db.fetchall(
|
|
"SELECT * FROM batches WHERE trade_date = ? AND state = 'failed'",
|
|
(TRADE_DATE,),
|
|
)
|
|
self.assertTrue(any("release group not switched" in str(row["error"] or "") for row in abandoned))
|
|
audit = self.db.fetchone(
|
|
"SELECT * FROM audit_log WHERE action = 'release-group' ORDER BY id DESC"
|
|
)
|
|
self.assertIn("valuation", str(audit["detail"]))
|
|
# still missing → evening retries keep trying
|
|
self.assertIn("daily", self.pipe.missing_official_datasets(TRADE_DATE))
|
|
|
|
def test_failure_keeps_previous_complete_version_serving(self) -> None:
|
|
first = self.pipe.run_dataset("daily", TRADE_DATE)
|
|
self.transport.empty = {"daily_basic"}
|
|
results = self.pipe.run_eod_missing(TRADE_DATE)
|
|
# incomplete A-group restages daily with the others; valuation fails → no A switch
|
|
self.assertEqual(results["daily"]["state"], "failed")
|
|
self.assertEqual(results["valuation"]["state"], "failed")
|
|
# the already-published daily batch is untouched and keeps serving
|
|
self.assertEqual(self.pipe.active_batch("daily", TRADE_DATE), first["batch_id"])
|
|
pubs = publications_map(self.db, TRADE_DATE)
|
|
self.assertEqual(pubs["daily"], first["batch_id"])
|
|
self.assertNotIn("valuation", pubs)
|
|
self.assertNotIn("moneyflow", pubs)
|
|
self.assertNotIn("auction", pubs)
|
|
# B-group is an independent boundary and may still publish
|
|
self.assertEqual(results["index_daily"]["state"], "published")
|
|
payload = V1API(self.db, self.pipe, self.pipe.settings).handle(
|
|
"/v1/bars/daily", {"date": [TRADE_DATE], "code": ["600000.SH"]}
|
|
)
|
|
self.assertEqual(payload["meta"]["batch_id"], first["batch_id"])
|
|
|
|
def test_partial_group_retry_does_not_mix_batches(self) -> None:
|
|
"""Already-published A members must be restaged with missing ones."""
|
|
first_daily = self.pipe.run_dataset("daily", TRADE_DATE)
|
|
first_moneyflow = self.pipe.run_dataset("moneyflow", TRADE_DATE)
|
|
results = self.pipe.run_eod_missing(TRADE_DATE)
|
|
# A-group switched as one boundary; B-group (index) also published
|
|
for name in (*GROUP_A, "stocks"):
|
|
self.assertEqual(results[name]["state"], "published", name)
|
|
self.assertEqual(results["index_daily"]["state"], "published")
|
|
pubs = self.db.fetchall(
|
|
"SELECT dataset, active_batch, published_at FROM publications WHERE trade_date = ?",
|
|
(TRADE_DATE,),
|
|
)
|
|
by_ds = {str(row["dataset"]): row for row in pubs}
|
|
# old partial batches replaced — no cross-batch mix of the first wave
|
|
self.assertNotEqual(by_ds["daily"]["active_batch"], first_daily["batch_id"])
|
|
self.assertNotEqual(by_ds["moneyflow"]["active_batch"], first_moneyflow["batch_id"])
|
|
a_times = {by_ds[name]["published_at"] for name in (*GROUP_A, "stocks")}
|
|
self.assertEqual(len(a_times), 1)
|
|
# serving resolves the new complete A-group batches
|
|
api = V1API(self.db, self.pipe, self.pipe.settings)
|
|
daily = api.handle("/v1/bars/daily", {"date": [TRADE_DATE], "code": ["600000.SH"]})
|
|
self.assertEqual(daily["meta"]["batch_id"], results["daily"]["batch_id"])
|
|
self.assertEqual(daily["meta"]["batch_id"], by_ds["daily"]["active_batch"])
|
|
|
|
def test_reads_during_switch_see_old_state_until_commit(self) -> None:
|
|
snapshots: list[dict] = []
|
|
|
|
def watcher() -> None:
|
|
with self.db.connect() as connection:
|
|
rows = connection.execute(
|
|
"SELECT dataset, active_batch FROM publications WHERE trade_date = ?",
|
|
(TRADE_DATE,),
|
|
).fetchall()
|
|
snapshots.append({str(row["dataset"]): row["active_batch"] for row in rows})
|
|
|
|
self.pipe.before_commit = watcher
|
|
self.pipe.run_eod_batch_a(TRADE_DATE)
|
|
# inside the switch transaction the group was still invisible
|
|
self.assertEqual(snapshots[0], {})
|
|
after = publications_map(self.db, TRADE_DATE)
|
|
self.assertEqual(set(after), {*GROUP_A, "stocks"})
|
|
|
|
def test_switch_crash_rolls_back_whole_group(self) -> None:
|
|
def explode() -> None:
|
|
raise RuntimeError("killed mid-switch")
|
|
|
|
self.pipe.before_commit = explode
|
|
with self.assertRaises(RuntimeError):
|
|
self.pipe.run_eod_batch_a(TRADE_DATE)
|
|
self.assertEqual(publications_map(self.db, TRADE_DATE), {})
|
|
for table in ("eod_bars", "eod_valuation", "eod_moneyflow", "eod_auction", "eod_stocks"):
|
|
rows = self.db.fetchall(f"SELECT * FROM {table} WHERE trade_date = ?", (TRADE_DATE,))
|
|
self.assertEqual(rows, [], table)
|
|
|
|
def test_duplicate_runs_are_idempotent(self) -> None:
|
|
self.pipe.run_eod_batch_a(TRADE_DATE)
|
|
self.pipe.run_eod_batch_b(TRADE_DATE)
|
|
batches_before = {
|
|
str(row["batch_id"])
|
|
for row in self.db.fetchall("SELECT batch_id FROM batches WHERE trade_date = ?", (TRADE_DATE,))
|
|
}
|
|
calls_before = len(self.transport.calls)
|
|
again = self.pipe.run_eod_missing(TRADE_DATE)
|
|
self.assertEqual({item["state"] for item in again.values()}, {"skipped"})
|
|
self.assertEqual({item["reason"] for item in again.values()}, {"already_published"})
|
|
batches_after = {
|
|
str(row["batch_id"])
|
|
for row in self.db.fetchall("SELECT batch_id FROM batches WHERE trade_date = ?", (TRADE_DATE,))
|
|
}
|
|
self.assertEqual(batches_after, batches_before)
|
|
self.assertEqual(len(self.transport.calls), calls_before)
|
|
self.assertEqual(self.pipe.missing_official_datasets(TRADE_DATE), [])
|
|
|
|
def test_cross_gate_failure_blocks_switch(self) -> None:
|
|
transport = GroupTransport()
|
|
pipe, db = make_pipe(
|
|
transport,
|
|
quality_extra={"cross_gates": [
|
|
{"left": "daily", "right": "moneyflow", "min_key_overlap": 1.0},
|
|
]},
|
|
)
|
|
pipe.ingest_reference(TRADE_DATE)
|
|
transport.keep_rows["moneyflow"] = 1 # moneyflow covers only half the market
|
|
results = pipe.run_eod_batch_a(TRADE_DATE)
|
|
self.assertEqual(results["moneyflow"]["state"], "failed")
|
|
self.assertIn("cross gate", str(results["moneyflow"]["error"]))
|
|
self.assertEqual(publications_map(db, TRADE_DATE), {})
|
|
|
|
def test_stocks_master_and_snapshot_switch_together_or_not_at_all(self) -> None:
|
|
original = [
|
|
{"ts_code": "600000.SH", "symbol": "600000", "name": "浦发银行", "area": "上海",
|
|
"industry": "银行", "market": "主板", "list_status": "L", "list_date": "19991110"},
|
|
{"ts_code": "920071.BJ", "symbol": "920071", "name": "N金钛", "area": "辽宁",
|
|
"industry": "小金属", "market": "北交所", "list_status": "L", "list_date": "20240901"},
|
|
]
|
|
renamed = [dict(original[0]), {**original[1], "name": "金钛股份"}]
|
|
self.transport.stocks = renamed
|
|
self.pipe.run_eod_batch_a(TRADE_DATE)
|
|
master = self.db.fetchone("SELECT name FROM stock_master WHERE ts_code = '920071.BJ'")
|
|
self.assertEqual(master["name"], "金钛股份")
|
|
stocks_pub = self.db.fetchone(
|
|
"SELECT active_batch FROM publications WHERE dataset = 'stocks' AND trade_date = ?",
|
|
(TRADE_DATE,),
|
|
)
|
|
self.assertIsNotNone(stocks_pub)
|
|
|
|
# failure path: rename staged but the group is blocked → master stays untouched
|
|
transport = GroupTransport()
|
|
transport.stocks = original
|
|
pipe, db = make_pipe(
|
|
transport,
|
|
quality_extra={"cross_gates": [
|
|
{"left": "daily", "right": "moneyflow", "min_key_overlap": 1.0},
|
|
]},
|
|
)
|
|
pipe.ingest_reference(TRADE_DATE) # master seeded with "N金钛"
|
|
transport.stocks = renamed
|
|
transport.keep_rows["moneyflow"] = 1
|
|
results = pipe.run_eod_batch_a(TRADE_DATE)
|
|
self.assertEqual(results["stocks"]["state"], "failed")
|
|
master = db.fetchone("SELECT name FROM stock_master WHERE ts_code = '920071.BJ'")
|
|
self.assertEqual(master["name"], "N金钛") # rename not applied
|
|
stocks_pub = db.fetchone(
|
|
"SELECT active_batch FROM publications WHERE dataset = 'stocks' AND trade_date = ?",
|
|
(TRADE_DATE,),
|
|
)
|
|
self.assertIsNone(stocks_pub)
|
|
|
|
|
|
class StocksRefreshAtomicTests(unittest.TestCase):
|
|
def setUp(self) -> None:
|
|
self.transport = GroupTransport()
|
|
self.pipe, self.db = make_pipe(self.transport)
|
|
self.pipe.ingest_reference(TRADE_DATE)
|
|
self.transport.stocks = [
|
|
{"ts_code": "600000.SH", "symbol": "600000", "name": "浦发银行", "area": "上海",
|
|
"industry": "银行", "market": "主板", "list_status": "L", "list_date": "19991110"},
|
|
{"ts_code": "920071.BJ", "symbol": "920071", "name": "N金钛", "area": "辽宁",
|
|
"industry": "小金属", "market": "北交所", "list_status": "L", "list_date": "20240901"},
|
|
]
|
|
first = self.pipe.refresh_stocks(TRADE_DATE)
|
|
self.assertEqual(first["state"], "published")
|
|
self.first_batch = first["batch_id"]
|
|
|
|
def test_refresh_keeps_master_when_snapshot_publish_fails(self) -> None:
|
|
self.transport.stocks = [
|
|
{"ts_code": "600000.SH", "symbol": "600000", "name": "浦发银行", "area": "上海",
|
|
"industry": "银行", "market": "主板", "list_status": "L", "list_date": "19991110"},
|
|
{"ts_code": "920071.BJ", "symbol": "920071", "name": "金钛股份", "area": "辽宁",
|
|
"industry": "小金属", "market": "北交所", "list_status": "L", "list_date": "20240901"},
|
|
]
|
|
|
|
def explode() -> None:
|
|
raise RuntimeError("snapshot switch killed")
|
|
|
|
self.pipe.before_commit = explode
|
|
with self.assertRaises(RuntimeError):
|
|
self.pipe.refresh_stocks(TRADE_DATE)
|
|
master = self.db.fetchone("SELECT name FROM stock_master WHERE ts_code = '920071.BJ'")
|
|
self.assertEqual(master["name"], "N金钛") # rename not applied
|
|
self.assertEqual(self.pipe.active_batch("stocks", TRADE_DATE), self.first_batch)
|
|
audit = self.db.fetchone(
|
|
"SELECT * FROM audit_log WHERE action = 'stocks-refresh' ORDER BY id DESC"
|
|
)
|
|
self.assertIn("failed", str(audit["detail"]))
|
|
self.assertIn("snapshot switch killed", str(audit["detail"]))
|
|
|
|
def test_refresh_keeps_master_when_quality_gate_rejects(self) -> None:
|
|
self.transport.stocks = [] # empty → hard fail before publish
|
|
with self.assertRaises(Exception):
|
|
self.pipe.refresh_stocks(TRADE_DATE)
|
|
master = self.db.fetchone("SELECT name FROM stock_master WHERE ts_code = '920071.BJ'")
|
|
self.assertEqual(master["name"], "N金钛")
|
|
self.assertEqual(self.pipe.active_batch("stocks", TRADE_DATE), self.first_batch)
|
|
audit = self.db.fetchone(
|
|
"SELECT * FROM audit_log WHERE action = 'stocks-refresh' ORDER BY id DESC"
|
|
)
|
|
self.assertIn("failed", str(audit["detail"]))
|
|
|
|
|
|
class ForceBoundaryEntryTests(unittest.TestCase):
|
|
"""CLI force / admin backfill must rebuild the full A/B boundary."""
|
|
|
|
def setUp(self) -> None:
|
|
self.transport = GroupTransport()
|
|
self.pipe, self.db = make_pipe(self.transport)
|
|
self.pipe.ingest_reference(TRADE_DATE)
|
|
self.first = self.pipe.run_eod_batch_a(TRADE_DATE)
|
|
self.pipe.run_eod_batch_b(TRADE_DATE)
|
|
|
|
def test_force_republish_valuation_rebuilds_whole_a_group(self) -> None:
|
|
before = publications_map(self.db, TRADE_DATE)
|
|
results = self.pipe.force_republish_boundary("valuation", TRADE_DATE)
|
|
self.assertEqual({item["state"] for item in results.values()}, {"published"})
|
|
after = publications_map(self.db, TRADE_DATE)
|
|
for name in (*GROUP_A, "stocks"):
|
|
self.assertNotEqual(after[name], before[name], name)
|
|
self.assertEqual(after[name], results[name]["batch_id"], name)
|
|
# B-group left alone
|
|
self.assertEqual(after["index_daily"], before["index_daily"])
|
|
pubs = self.db.fetchall(
|
|
"SELECT dataset, published_at FROM publications WHERE trade_date = ?",
|
|
(TRADE_DATE,),
|
|
)
|
|
a_times = {row["published_at"] for row in pubs if row["dataset"] in {*GROUP_A, "stocks"}}
|
|
self.assertEqual(len(a_times), 1)
|
|
|
|
def test_force_republish_index_rebuilds_only_b_group(self) -> None:
|
|
before = publications_map(self.db, TRADE_DATE)
|
|
results = self.pipe.force_republish_boundary("index_daily", TRADE_DATE)
|
|
self.assertEqual(results["index_daily"]["state"], "published")
|
|
after = publications_map(self.db, TRADE_DATE)
|
|
self.assertNotEqual(after["index_daily"], before["index_daily"])
|
|
for name in GROUP_A:
|
|
self.assertEqual(after[name], before[name], name)
|
|
|
|
def test_admin_backfill_official_dataset_uses_boundary(self) -> None:
|
|
from datahub.admin_api import AdminAPI
|
|
from datahub.auth import AuthService
|
|
from datahub.crypto import SecretVault
|
|
from datahub.scheduler import Scheduler
|
|
from datahub.serving import ApiError
|
|
|
|
vault = SecretVault(self.pipe.settings.encryption_key)
|
|
auth = AuthService(self.db, vault, self.pipe.settings.api_token, "StartPass1")
|
|
admin = AdminAPI(self.db, self.pipe, Scheduler(self.db, self.pipe), auth)
|
|
before = publications_map(self.db, TRADE_DATE)
|
|
result = admin.backfill("moneyflow", TRADE_DATE, "StartPass1", f"moneyflow:{TRADE_DATE}", "tester")
|
|
self.assertEqual(result["moneyflow"]["state"], "published")
|
|
after = publications_map(self.db, TRADE_DATE)
|
|
for name in (*GROUP_A, "stocks"):
|
|
self.assertNotEqual(after[name], before[name], name)
|
|
# bad password / wrong confirm still rejected
|
|
with self.assertRaises(ApiError):
|
|
admin.backfill("daily", TRADE_DATE, "wrong", f"daily:{TRADE_DATE}", "tester")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|