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, QualityError from datahub.scheduler import Scheduler from datahub.settings import Settings from tests.fixtures import TRADE_DATE, fake_transport from tests.test_eod_retry import clock_at FIELD_GATES = { "valuation": { "fields": [ "turnover_rate", "volume_ratio", "total_mv", "circ_mv", "pe_ttm", "pb", "ps_ttm", "dv_ttm", ], "min_nonnull_rate": 0.9, "min_nonnull_rate_by_field": {"pe_ttm": 0.5, "dv_ttm": 0.3}, "max_nonnull_drop_vs_prev": 0.15, "max_nonfinite_rate": 0.01, }, } class ValuationTransport: """fake_transport with switchable daily_basic degradation modes.""" def __init__(self) -> None: self.mode = "ok" def __call__(self, api_name: str, params: dict, fields: str): rows = fake_transport(api_name, params, fields) if api_name != "daily_basic": return rows trade_date = str(params.get("trade_date") or "") if trade_date: rows = [{**row, "trade_date": trade_date} for row in rows] if self.mode == "ok": return rows patched = [] for row in rows: item = dict(row) if self.mode == "fields_all_null": item["volume_ratio"] = None item["dv_ttm"] = None elif self.mode == "vr_all_null": item["volume_ratio"] = None elif self.mode == "dv_all_null": item["dv_ttm"] = None elif self.mode == "nonfinite": item["volume_ratio"] = float("inf") patched.append(item) return patched def make_pipe(transport, quality_extra=None, clock=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": 3, "publication_generations": 3, "job_run_retain_days": 90, "staging_retain_days": 14, "field_gates": FIELD_GATES, } 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, clock=clock) pipe._tmp = tmp return pipe, db class ValuationFieldGateTests(unittest.TestCase): def setUp(self) -> None: self.transport = ValuationTransport() self.pipe, self.db = make_pipe(self.transport) self.pipe.ingest_reference(TRADE_DATE) def _active(self) -> str | None: row = self.db.fetchone( "SELECT active_batch FROM publications WHERE dataset='valuation' AND trade_date=?", (TRADE_DATE,), ) return str(row["active_batch"]) if row else None def test_normal_batch_with_legit_dv_nulls_passes(self) -> None: result = self.pipe.run_dataset("valuation", TRADE_DATE) self.assertEqual(result["state"], "published") fields = result["quality"]["fields"] # fixture: 1 of 2 stocks has null dv_ttm → 0.5 non-null ≥ 0.3 floor self.assertEqual(fields["dv_ttm"]["nonnull_rate"], 0.5) self.assertEqual(fields["volume_ratio"]["nonnull_rate"], 1.0) self.assertFalse(result["quality"]["errors"]) def test_all_null_fields_rejected_and_prev_batch_kept(self) -> None: first = self.pipe.run_dataset("valuation", TRADE_DATE) self.transport.mode = "fields_all_null" with self.assertRaises(QualityError) as ctx: self.pipe.run_dataset("valuation", TRADE_DATE) errors = "; ".join(ctx.exception.report["errors"]) self.assertIn("field gate: valuation.volume_ratio non-null rate 0.0000 < 0.9", errors) self.assertIn("field gate: valuation.dv_ttm non-null rate 0.0000 < 0.3", errors) # previous good publication stays active self.assertEqual(self._active(), first["batch_id"]) # rejected batch left staged with readable error + field stats rejected = self.db.fetchone( "SELECT * FROM batches WHERE state='staged' AND dataset='valuation' ORDER BY started_at DESC", ) self.assertIsNotNone(rejected) self.assertIn("field gate: valuation.volume_ratio", str(rejected["error"])) import json quality = json.loads(rejected["quality_json"]) self.assertEqual(quality["fields"]["volume_ratio"]["nonnull"], 0) self.assertEqual(quality["fields"]["dv_ttm"]["nonnull"], 0) def test_volume_ratio_all_null_alone_rejected(self) -> None: self.pipe.run_dataset("valuation", TRADE_DATE) self.transport.mode = "vr_all_null" with self.assertRaises(QualityError): self.pipe.run_dataset("valuation", TRADE_DATE) self.assertEqual( self.db.fetchone( "SELECT active_batch FROM publications WHERE dataset='valuation' AND trade_date=?", (TRADE_DATE,), )["active_batch"], "20240902-valuation-001", ) def test_dv_ttm_all_null_rejected_by_floor_and_collapse(self) -> None: prev_day = "20240830" prev = self.pipe.run_dataset("valuation", prev_day) # prev dv nonnull 0.5 self.transport.mode = "dv_all_null" with self.assertRaises(QualityError) as ctx: self.pipe.run_dataset("valuation", TRADE_DATE) errors = "; ".join(ctx.exception.report["errors"]) self.assertIn("field gate: valuation.dv_ttm non-null rate 0.0000 < 0.3", errors) self.assertIn(f"dropped > 0.15 vs prev batch {prev['batch_id']}", errors) def test_nonfinite_values_rejected(self) -> None: self.pipe.run_dataset("valuation", TRADE_DATE) rows = self.pipe.adapter.normalize( "valuation", self.pipe._guarded_fetch("valuation", {"trade_date": TRADE_DATE}) ) for row in rows: row["volume_ratio"] = float("inf") with self.assertRaises(QualityError) as ctx: self.pipe.run_dataset("valuation", TRADE_DATE, prepared_rows=rows) errors = "; ".join(ctx.exception.report["errors"]) self.assertIn("field gate: valuation.volume_ratio non-finite rate 1.0000 > 0.01", errors) def test_gate_off_when_not_configured(self) -> None: pipe, _db = make_pipe(ValuationTransport(), quality_extra={"field_gates": {}}) pipe.ingest_reference(TRADE_DATE) pipe.adapter._transport.mode = "fields_all_null" result = pipe.run_dataset("valuation", TRADE_DATE) self.assertEqual(result["state"], "published") # legacy behavior when unconfigured def test_gate_applies_to_any_configured_dataset(self) -> None: gates = {"daily": {"fields": ["volume"], "min_nonnull_rate": 0.9, "max_nonfinite_rate": 0.01}} pipe, _db = make_pipe(ValuationTransport(), quality_extra={"field_gates": gates}) pipe.ingest_reference(TRADE_DATE) def null_volume(api_name, params, fields): if api_name != "daily": return fake_transport(api_name, params, fields) rows = fake_transport(api_name, params, fields) for row in rows: row["vol"] = None return rows pipe.adapter._transport = null_volume with self.assertRaises(QualityError) as ctx: pipe.run_dataset("daily", TRADE_DATE) errors = "; ".join(ctx.exception.report["errors"]) self.assertIn("field gate: daily.volume non-null rate 0.0000 < 0.9", errors) class GateRetryInterplayTests(unittest.TestCase): def test_rejected_valuation_stays_missing_and_retry_publishes_later(self) -> None: transport = ValuationTransport() transport.mode = "fields_all_null" tmp = tempfile.TemporaryDirectory() self.addCleanup(tmp.cleanup) db = HubDB(Path(tmp.name) / "hub.db") adapter = TushareAdapter("x", transport=transport) settings = Settings( encryption_key=SecretVault.generate_key(), db_path=db.path, backup_dir=Path(tmp.name) / "backups", quality={"field_gates": FIELD_GATES, "max_publish_attempts": 2}, ) pipe = Pipeline(db, adapter, settings) pipe.ingest_reference("20240902") sched = Scheduler(db, pipe) sched.tick(clock_at("20240902", 15, 5)) # valuation rejected by field gate sched.tick(clock_at("20240902", 15, 10)) self.assertIn("valuation", pipe.missing_official_datasets("20240902")) self.assertEqual( pipe.active_batch("valuation", "20240902"), None, ) transport.mode = "ok" sched.tick(clock_at("20240902", 15, 45)) # retry passes the gate self.assertNotIn("valuation", pipe.missing_official_datasets("20240902")) rows = db.fetchall("SELECT * FROM eod_valuation WHERE trade_date='20240902'") self.assertTrue(rows) self.assertTrue(all(row["volume_ratio"] is not None for row in rows)) class ForceRepublishTests(unittest.TestCase): def test_run_dataset_over_published_keeps_prev_for_rollback(self) -> None: transport = ValuationTransport() pipe, db = make_pipe(transport) pipe.ingest_reference(TRADE_DATE) first = pipe.run_dataset("valuation", TRADE_DATE) transport.mode = "vr_all_null" with self.assertRaises(QualityError): pipe.run_dataset("valuation", TRADE_DATE) # gate holds: bad re-publish refused transport.mode = "ok" second = pipe.run_dataset("valuation", TRADE_DATE) # CLI --force path self.assertNotEqual(first["batch_id"], second["batch_id"]) pub = db.fetchone( "SELECT * FROM publications WHERE dataset='valuation' AND trade_date=?", (TRADE_DATE,), ) self.assertEqual(pub["active_batch"], second["batch_id"]) self.assertEqual(pub["prev_batch"], first["batch_id"]) rolled = pipe.rollback("valuation", TRADE_DATE, actor="cli") self.assertEqual(rolled["active_batch"], first["batch_id"]) if __name__ == "__main__": unittest.main()