from __future__ import annotations import unittest from datetime import date, timedelta 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.serving import ApiError, V1API from datahub.settings import Settings from tests.fixtures import fake_transport OPEN_DATES = ["20240826", "20240827", "20240828", "20240829", "20240830", "20240902", "20240903"] EMPTY_UPSTREAM = {"20240828"} # one date the upstream cannot serve def build_calendar(open_dates: list[str], span_days: int = 16) -> list[dict]: start = date(int(open_dates[0][:4]), int(open_dates[0][4:6]), int(open_dates[0][6:8])) rows = [] open_set = set(open_dates) for offset in range(span_days): cursor = start + timedelta(days=offset) compact = cursor.strftime("%Y%m%d") rows.append( {"exchange": "SSE", "cal_date": compact, "is_open": 1 if compact in open_set else 0, "pretrade_date": compact} ) return rows def moneyflow_rows(day: str) -> list[dict]: return [ { "ts_code": "600000.SH", "trade_date": day, "buy_sm_amount": 10 + int(day[-2:]), "sell_sm_amount": 8, "buy_md_amount": 20, "sell_md_amount": 15, "buy_lg_amount": 30, "sell_lg_amount": 25, "buy_elg_amount": 40, "sell_elg_amount": 35, "net_mf_amount": 17, }, { "ts_code": "000001.SZ", "trade_date": day, "buy_sm_amount": 11, "sell_sm_amount": 9, "buy_md_amount": 21, "sell_md_amount": 16, "buy_lg_amount": 31, "sell_lg_amount": 26, "buy_elg_amount": 41, "sell_elg_amount": 36, "net_mf_amount": 18, }, ] class MoneyflowHistoryTransport: def __init__(self) -> None: self.calendar = build_calendar(OPEN_DATES) self.moneyflow_fetches: list[str] = [] def __call__(self, api_name: str, params: dict, fields: str): if api_name == "trade_cal": start = str(params.get("start_date") or "") end = str(params.get("end_date") or "99999999") return [row for row in self.calendar if start <= row["cal_date"] <= end] if api_name == "moneyflow": day = str(params.get("trade_date") or "") self.moneyflow_fetches.append(day) if day in EMPTY_UPSTREAM: return [] return moneyflow_rows(day) return fake_transport(api_name, params, fields) class MoneyflowBackfillTests(unittest.TestCase): def setUp(self) -> None: self.transport = MoneyflowHistoryTransport() tmp = tempfile.TemporaryDirectory() self.addCleanup(tmp.cleanup) self.db = HubDB(Path(tmp.name) / "hub.db") adapter = TushareAdapter("x", transport=self.transport) settings = Settings( encryption_key=SecretVault.generate_key(), db_path=self.db.path, backup_dir=Path(tmp.name) / "backups", quality={"max_publish_attempts": 2, "publication_generations": 3}, ) self.pipe = Pipeline(self.db, adapter, settings) self.pipe.ingest_reference("20240903") self.api = V1API(self.db, self.pipe, settings) def test_backfill_publishes_window_and_reports_failures(self) -> None: result = self.pipe.backfill_moneyflow_history(end_date="20240903", trading_days=5) published = [item["trade_date"] for item in result["published"]] self.assertEqual(published, ["20240829", "20240830", "20240902", "20240903"]) self.assertEqual(result["failed"][0]["trade_date"], "20240828") self.assertFalse(result["ok"]) rows = self.db.fetchall("SELECT * FROM eod_moneyflow WHERE trade_date='20240830'") self.assertEqual(len(rows), 2) def test_backfill_is_idempotent(self) -> None: self.pipe.backfill_moneyflow_history(end_date="20240903", trading_days=5) fetches_after_first = list(self.transport.moneyflow_fetches) second = self.pipe.backfill_moneyflow_history(end_date="20240903", trading_days=5) # only the still-missing date is re-fetched; published dates are skipped self.assertEqual(self.transport.moneyflow_fetches[len(fetches_after_first):], ["20240828"]) self.assertEqual(len(second["skipped"]), 4) def test_point_query_on_backfilled_date_serves_data(self) -> None: self.pipe.backfill_moneyflow_history(end_date="20240903", trading_days=5) payload = self.api.handle("/v1/moneyflow", {"date": ["20240830"]}) self.assertEqual(len(payload["data"]), 2) self.assertEqual(payload["data"][0]["net_mf_amount"], 180000.0) def test_unpublished_point_below_window_is_identifiable(self) -> None: self.pipe.backfill_moneyflow_history(end_date="20240903", trading_days=5) with self.assertRaises(ApiError) as ctx: self.api.handle("/v1/moneyflow", {"date": ["20240801"]}) extra = ctx.exception.extra self.assertEqual(extra["available_from"], "20240829") # window starts at the first published date self.assertEqual(extra["available_to"], "20240903") self.assertEqual(extra["reason"], "history_not_backfilled") self.assertEqual(extra["expected_at"], "15:05+08:00") def test_range_query_flags_missing_dates(self) -> None: self.pipe.backfill_moneyflow_history(end_date="20240903", trading_days=5) payload = self.api.handle("/v1/moneyflow", {"from": ["20240828"], "to": ["20240903"]}) coverage = payload["meta"]["coverage"] self.assertFalse(coverage["complete"]) self.assertEqual(coverage["missing_count"], 1) self.assertEqual(coverage["missing_sample"], ["20240828"]) self.assertTrue(payload["meta"]["incomplete"]) if __name__ == "__main__": unittest.main()