feat: operationalize local Tushare Qlib research
This commit is contained in:
+324
-6
@@ -13,9 +13,11 @@ sys.path.insert(0, str(ROOT))
|
||||
from adapters.tushare_local import (
|
||||
TushareMirrorError,
|
||||
build_qlib_provider,
|
||||
inventory_mirror,
|
||||
tushare_to_qlib_symbol,
|
||||
verify_qlib_provider,
|
||||
)
|
||||
from tools.tushare_qlib import _assert_json_outside_provider
|
||||
|
||||
|
||||
class TushareSymbolTest(unittest.TestCase):
|
||||
@@ -26,6 +28,20 @@ class TushareSymbolTest(unittest.TestCase):
|
||||
with self.assertRaises(ValueError):
|
||||
tushare_to_qlib_symbol("T600018.SH")
|
||||
|
||||
def test_provider_json_output_cannot_mutate_provider_tree(self):
|
||||
with self.assertRaisesRegex(
|
||||
ValueError,
|
||||
"outside the immutable provider",
|
||||
):
|
||||
_assert_json_outside_provider(
|
||||
"/tmp/provider/evidence.json",
|
||||
"/tmp/provider",
|
||||
)
|
||||
_assert_json_outside_provider(
|
||||
"/tmp/evidence.json",
|
||||
"/tmp/provider",
|
||||
)
|
||||
|
||||
|
||||
@unittest.skipUnless(
|
||||
importlib.util.find_spec("pandas") is not None
|
||||
@@ -33,7 +49,14 @@ class TushareSymbolTest(unittest.TestCase):
|
||||
"pandas and pyarrow are optional outside the Qlib research environment",
|
||||
)
|
||||
class TushareProviderTest(unittest.TestCase):
|
||||
def _mirror(self, root: Path, *, bad_ohlc: bool = False) -> Path:
|
||||
def _mirror(
|
||||
self,
|
||||
root: Path,
|
||||
*,
|
||||
bad_ohlc: bool = False,
|
||||
advanced: bool = False,
|
||||
shadowed_daily: bool = False,
|
||||
) -> Path:
|
||||
import pandas
|
||||
|
||||
mirror = root / "mirror"
|
||||
@@ -43,7 +66,10 @@ class TushareProviderTest(unittest.TestCase):
|
||||
(parquet_root / "stock_basic").mkdir(parents=True)
|
||||
dates = pandas.bdate_range("2024-01-02", periods=8)
|
||||
daily_rows = []
|
||||
for symbol, base in (("600000.SH", 10.0), ("000001.SZ", 20.0)):
|
||||
symbols = [("600000.SH", 10.0), ("000001.SZ", 20.0)]
|
||||
if advanced:
|
||||
symbols.append(("000002.SZ", 30.0))
|
||||
for symbol, base in symbols:
|
||||
previous = base
|
||||
for index, value in enumerate(dates):
|
||||
close = base * (1.01 ** index)
|
||||
@@ -78,11 +104,121 @@ class TushareProviderTest(unittest.TestCase):
|
||||
),
|
||||
"stock_basic/list_status=L.parquet": pandas.DataFrame(
|
||||
{
|
||||
"ts_code": ["600000.SH", "000001.SZ"],
|
||||
"curr_type": ["CNY", "CNY"],
|
||||
"ts_code": [value[0] for value in symbols],
|
||||
"curr_type": ["CNY"] * len(symbols),
|
||||
}
|
||||
),
|
||||
}
|
||||
partitions = {}
|
||||
if advanced:
|
||||
adj_rows = []
|
||||
for symbol, _ in symbols:
|
||||
for index, value in enumerate(dates):
|
||||
factor = 2.0 if symbol == "600000.SH" and index >= 4 else 1.0
|
||||
adj_rows.append(
|
||||
{
|
||||
"ts_code": symbol,
|
||||
"trade_date": value.strftime("%Y%m%d"),
|
||||
"adj_factor": factor,
|
||||
}
|
||||
)
|
||||
index_rows = []
|
||||
previous = 1000.0
|
||||
for index, value in enumerate(dates):
|
||||
close = 1000.0 + index * 10.0
|
||||
index_rows.append(
|
||||
{
|
||||
"ts_code": "000905.SH",
|
||||
"trade_date": value.strftime("%Y%m%d"),
|
||||
"open": close - 2.0,
|
||||
"high": close + 5.0,
|
||||
"low": close - 5.0,
|
||||
"close": close,
|
||||
"pre_close": previous,
|
||||
"change": close - previous,
|
||||
"pct_chg": (close / previous - 1.0) * 100.0,
|
||||
"vol": 2000.0,
|
||||
"amount": close * 200.0,
|
||||
}
|
||||
)
|
||||
previous = close
|
||||
first_snapshot = dates[0].strftime("%Y%m%d")
|
||||
second_snapshot = dates[4].strftime("%Y%m%d")
|
||||
anchor_weight_rows = [
|
||||
{
|
||||
"index_code": "000905.SH",
|
||||
"con_code": "600000.SH",
|
||||
"trade_date": "20231229",
|
||||
"weight": 50.0,
|
||||
},
|
||||
{
|
||||
"index_code": "000905.SH",
|
||||
"con_code": "000001.SZ",
|
||||
"trade_date": "20231229",
|
||||
"weight": 50.0,
|
||||
},
|
||||
]
|
||||
weight_rows = [
|
||||
{
|
||||
"index_code": "000905.SH",
|
||||
"con_code": "600000.SH",
|
||||
"trade_date": first_snapshot,
|
||||
"weight": 50.0,
|
||||
},
|
||||
{
|
||||
"index_code": "000905.SH",
|
||||
"con_code": "000001.SZ",
|
||||
"trade_date": first_snapshot,
|
||||
"weight": 50.0,
|
||||
},
|
||||
{
|
||||
"index_code": "000905.SH",
|
||||
"con_code": "000001.SZ",
|
||||
"trade_date": second_snapshot,
|
||||
"weight": 55.0,
|
||||
},
|
||||
{
|
||||
"index_code": "000905.SH",
|
||||
"con_code": "000002.SZ",
|
||||
"trade_date": second_snapshot,
|
||||
"weight": 45.0,
|
||||
},
|
||||
]
|
||||
frames.update(
|
||||
{
|
||||
"adj_factor/month=2024-01.parquet": pandas.DataFrame(
|
||||
adj_rows
|
||||
),
|
||||
(
|
||||
"index_daily/"
|
||||
"ts_code=000905_SH_year=2024.parquet"
|
||||
): pandas.DataFrame(index_rows),
|
||||
(
|
||||
"index_weight/"
|
||||
"index_code=000905_SH_year=2023.parquet"
|
||||
): pandas.DataFrame(anchor_weight_rows),
|
||||
(
|
||||
"index_weight/"
|
||||
"index_code=000905_SH_year=2024.parquet"
|
||||
): pandas.DataFrame(weight_rows),
|
||||
}
|
||||
)
|
||||
partitions.update(
|
||||
{
|
||||
(
|
||||
"index_daily/"
|
||||
"ts_code=000905_SH_year=2024.parquet"
|
||||
): "ts_code=000905_SH/year=2024",
|
||||
(
|
||||
"index_weight/"
|
||||
"index_code=000905_SH_year=2023.parquet"
|
||||
): "index_code=000905_SH/year=2023",
|
||||
(
|
||||
"index_weight/"
|
||||
"index_code=000905_SH_year=2024.parquet"
|
||||
): "index_code=000905_SH/year=2024",
|
||||
}
|
||||
)
|
||||
state = mirror / "data/state.sqlite3"
|
||||
connection = sqlite3.connect(state)
|
||||
connection.execute(
|
||||
@@ -117,13 +253,24 @@ class TushareProviderTest(unittest.TestCase):
|
||||
(
|
||||
index,
|
||||
api,
|
||||
path.stem,
|
||||
partitions.get(relative, path.stem),
|
||||
len(frame),
|
||||
path.stat().st_size,
|
||||
hashlib.sha256(path.read_bytes()).hexdigest(),
|
||||
str(path),
|
||||
),
|
||||
)
|
||||
if shadowed_daily:
|
||||
daily_path = parquet_root / "daily/month=2024-01.parquet"
|
||||
connection.execute(
|
||||
"""
|
||||
INSERT INTO jobs VALUES
|
||||
(0, 'daily', 'month=2024-01', 'completed', 1, 1,
|
||||
'obsolete-checksum', ?, NULL,
|
||||
'2024-01-01T00:00:00Z', NULL)
|
||||
""",
|
||||
(str(daily_path),),
|
||||
)
|
||||
connection.execute(
|
||||
"""
|
||||
INSERT INTO jobs VALUES
|
||||
@@ -150,7 +297,9 @@ class TushareProviderTest(unittest.TestCase):
|
||||
)
|
||||
self.assertFalse(result["investment_value_claim"])
|
||||
self.assertEqual(result["market"]["instrument_count"], 2)
|
||||
self.assertTrue(verify_qlib_provider(output)["ok"])
|
||||
verified = verify_qlib_provider(output)
|
||||
self.assertTrue(verified["ok"])
|
||||
self.assertTrue(verified["converter_source_matches_current"])
|
||||
manifest = json.loads(
|
||||
(output / "quant_os_tushare_manifest.json").read_text()
|
||||
)
|
||||
@@ -225,6 +374,175 @@ class TushareProviderTest(unittest.TestCase):
|
||||
self.assertAlmostEqual(high[1], 10.0)
|
||||
self.assertAlmostEqual(low[1], 8.0)
|
||||
|
||||
def test_inventory_ignores_shadowed_completed_file_version(self):
|
||||
with tempfile.TemporaryDirectory() as temporary:
|
||||
mirror = self._mirror(
|
||||
Path(temporary),
|
||||
shadowed_daily=True,
|
||||
)
|
||||
result = inventory_mirror(mirror)
|
||||
self.assertEqual(result["shadowed_completed_versions"], 1)
|
||||
self.assertEqual(
|
||||
result["shadowed_completed_versions_by_api"],
|
||||
{"daily": 1},
|
||||
)
|
||||
self.assertEqual(result["tables"]["daily"]["files"], 1)
|
||||
self.assertEqual(result["verified_completed_files"], 3)
|
||||
|
||||
def test_adjusted_real_benchmark_and_pit_universe(self):
|
||||
from array import array
|
||||
import math
|
||||
|
||||
with tempfile.TemporaryDirectory() as temporary:
|
||||
root = Path(temporary)
|
||||
mirror = self._mirror(root, advanced=True)
|
||||
output = root / "provider"
|
||||
result = build_qlib_provider(
|
||||
mirror,
|
||||
output,
|
||||
start="2024-01-02",
|
||||
end="2024-01-11",
|
||||
minimum_observations=999,
|
||||
benchmark_index_code="000905.SH",
|
||||
universe_index_code="000905.SH",
|
||||
expected_constituent_count=2,
|
||||
)
|
||||
self.assertEqual(result["schema_version"], 2)
|
||||
self.assertEqual(
|
||||
result["price_adjustment"]["mode"],
|
||||
"cumulative_adjusted_first_observation_anchor",
|
||||
)
|
||||
self.assertTrue(result["price_adjustment"]["component_complete"])
|
||||
self.assertFalse(result["price_adjustment"]["production_eligible"])
|
||||
self.assertEqual(
|
||||
result["benchmark"]["kind"],
|
||||
"tushare_index_daily",
|
||||
)
|
||||
self.assertTrue(result["market"]["point_in_time"])
|
||||
self.assertFalse(result["market"]["minimum_observations_applied"])
|
||||
self.assertFalse(result["market"]["same_session_membership_use"])
|
||||
self.assertEqual(
|
||||
result["market"]["availability_contract"],
|
||||
"conservative next-session effectiveness because the source "
|
||||
"does not expose a separately verified published_at timestamp",
|
||||
)
|
||||
self.assertEqual(result["market"]["instrument_count"], 3)
|
||||
self.assertTrue(
|
||||
result["source_stability"]["post_read_file_verification"]
|
||||
)
|
||||
self.assertEqual(
|
||||
{item["api_name"] for item in result["source_jobs"]},
|
||||
{
|
||||
"adj_factor",
|
||||
"daily",
|
||||
"index_daily",
|
||||
"index_weight",
|
||||
"stock_basic",
|
||||
"trade_cal",
|
||||
},
|
||||
)
|
||||
|
||||
factor = array("f")
|
||||
factor.frombytes(
|
||||
(output / "features/sh600000/factor.day.bin").read_bytes()
|
||||
)
|
||||
close = array("f")
|
||||
close.frombytes(
|
||||
(output / "features/sh600000/close.day.bin").read_bytes()
|
||||
)
|
||||
volume = array("f")
|
||||
volume.frombytes(
|
||||
(output / "features/sh600000/volume.day.bin").read_bytes()
|
||||
)
|
||||
change = array("f")
|
||||
change.frombytes(
|
||||
(output / "features/sh600000/change.day.bin").read_bytes()
|
||||
)
|
||||
self.assertAlmostEqual(factor[1], 1.0)
|
||||
self.assertAlmostEqual(factor[5], 2.0)
|
||||
raw_close = 10.0 * (1.01 ** 4)
|
||||
self.assertAlmostEqual(close[5] / factor[5], raw_close, places=5)
|
||||
self.assertAlmostEqual(volume[5], 50_000.0, places=3)
|
||||
self.assertTrue(math.isnan(change[1]))
|
||||
self.assertAlmostEqual(
|
||||
change[5],
|
||||
(raw_close * 2.0) / (10.0 * (1.01 ** 3)) - 1.0,
|
||||
places=5,
|
||||
)
|
||||
|
||||
index_close = array("f")
|
||||
index_close.frombytes(
|
||||
(output / "features/sh000905/close.day.bin").read_bytes()
|
||||
)
|
||||
self.assertAlmostEqual(index_close[1], 1000.0)
|
||||
market_lines = (
|
||||
output / "instruments/tushare_a.txt"
|
||||
).read_text(encoding="utf-8").splitlines()
|
||||
self.assertEqual(
|
||||
market_lines,
|
||||
[
|
||||
"SZ000001\t2024-01-02\t2024-01-11",
|
||||
"SZ000002\t2024-01-09\t2024-01-11",
|
||||
"SH600000\t2024-01-02\t2024-01-08",
|
||||
],
|
||||
)
|
||||
self.assertTrue(verify_qlib_provider(output)["ok"])
|
||||
|
||||
def test_pit_universe_fails_closed_on_missing_anchor_or_bad_roster(self):
|
||||
with tempfile.TemporaryDirectory() as temporary:
|
||||
root = Path(temporary)
|
||||
mirror = self._mirror(root, advanced=True)
|
||||
state = mirror / "data/state.sqlite3"
|
||||
connection = sqlite3.connect(state)
|
||||
connection.execute(
|
||||
"DELETE FROM jobs WHERE partition_key = ?",
|
||||
("index_code=000905_SH/year=2023",),
|
||||
)
|
||||
connection.commit()
|
||||
connection.close()
|
||||
with self.assertRaisesRegex(
|
||||
TushareMirrorError,
|
||||
"missing required partitions",
|
||||
):
|
||||
build_qlib_provider(
|
||||
mirror,
|
||||
root / "provider-missing-anchor",
|
||||
start="2024-01-02",
|
||||
end="2024-01-11",
|
||||
benchmark_index_code="000905.SH",
|
||||
universe_index_code="000905.SH",
|
||||
expected_constituent_count=2,
|
||||
)
|
||||
|
||||
with tempfile.TemporaryDirectory() as temporary:
|
||||
root = Path(temporary)
|
||||
mirror = self._mirror(root, advanced=True)
|
||||
with self.assertRaisesRegex(
|
||||
TushareMirrorError,
|
||||
"constituents; expected 3",
|
||||
):
|
||||
build_qlib_provider(
|
||||
mirror,
|
||||
root / "provider-bad-roster",
|
||||
start="2024-01-02",
|
||||
end="2024-01-11",
|
||||
benchmark_index_code="000905.SH",
|
||||
universe_index_code="000905.SH",
|
||||
expected_constituent_count=3,
|
||||
)
|
||||
with self.assertRaisesRegex(
|
||||
ValueError,
|
||||
"requires benchmark_index_code",
|
||||
):
|
||||
build_qlib_provider(
|
||||
mirror,
|
||||
root / "provider-synthetic-benchmark",
|
||||
start="2024-01-02",
|
||||
end="2024-01-11",
|
||||
universe_index_code="000905.SH",
|
||||
expected_constituent_count=2,
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
|
||||
@@ -76,6 +76,8 @@ class PublicStatusTests(unittest.TestCase):
|
||||
self.assertFalse(payload["marketScope"]["bseTradingEnabled"])
|
||||
self.assertIn("北交所", payload["marketScope"]["statement"])
|
||||
self.assertIn("fail closed", payload["marketScope"]["statement"])
|
||||
self.assertEqual(payload["summary"]["testsRecorded"], 325)
|
||||
self.assertEqual(payload["summary"]["optionalSkipsRecorded"], 1)
|
||||
|
||||
def test_public_payload_is_deliberately_non_secret_and_non_local(self):
|
||||
payload = _read_payload()
|
||||
@@ -99,6 +101,25 @@ class PublicStatusTests(unittest.TestCase):
|
||||
self.assertFalse(Path(evidence["path"]).is_absolute())
|
||||
self.assertNotIn("..", Path(evidence["path"]).parts)
|
||||
|
||||
def test_public_tushare_claims_cannot_regress_to_early_mirror(self):
|
||||
payload = _read_payload()
|
||||
encoded = json.dumps(payload, ensure_ascii=False, sort_keys=True)
|
||||
scorecard = (PROJECT / "gate_scorecard.json").read_text(
|
||||
encoding="utf-8"
|
||||
)
|
||||
stale_claims = (
|
||||
"93 个",
|
||||
"8.2%",
|
||||
"1993-10-29",
|
||||
"CSI500-2018-2025/de9657cd",
|
||||
"de5b96253ff765fa4e4be9736ddc2fad528f72c7721964313956d44b02969173",
|
||||
"3c550c0ba986f7eb736b421ac318fdc2a58eacaa65cfbdeaf6b3f920b5288304",
|
||||
"aed22a7b83e4defd45c4874e568d6aeafdabe43987cea301e0c1d09543dccb7c",
|
||||
)
|
||||
for stale_claim in stale_claims:
|
||||
self.assertNotIn(stale_claim, encoded)
|
||||
self.assertNotIn(stale_claim, scorecard)
|
||||
|
||||
def test_static_page_has_fail_closed_fallback_and_no_local_paths(self):
|
||||
page = INDEX_PATH.read_text(encoding="utf-8")
|
||||
self.assertIn('src="./status-data.js"', page)
|
||||
@@ -272,6 +293,10 @@ class PublicStatusTests(unittest.TestCase):
|
||||
evidence["joinquant-target-package-evidence"]["path"],
|
||||
"docs/JOINQUANT_TARGET_PACKAGE_EVIDENCE_2026-07-26.md",
|
||||
)
|
||||
self.assertEqual(
|
||||
evidence["tushare-csi500-research-run"]["path"],
|
||||
"evidence/tushare/CSI500-2018-2025/5bf19d2d/run.json",
|
||||
)
|
||||
|
||||
gaps = payload["deliveryGaps"]
|
||||
self.assertGreaterEqual(len(gaps["internal"]), 3)
|
||||
@@ -374,7 +399,7 @@ class PublicStatusTests(unittest.TestCase):
|
||||
def test_freshness_uses_daily_bars_not_the_later_trade_calendar(self):
|
||||
payload = _read_payload()
|
||||
freshness = {item["id"]: item for item in payload["freshness"]}
|
||||
self.assertEqual(freshness["project-audit"]["value"], "2026-07-30")
|
||||
self.assertEqual(freshness["project-audit"]["value"], "2026-07-31")
|
||||
self.assertEqual(
|
||||
freshness["joinquant-hosted-market-window"]["marketDataThrough"],
|
||||
"2024-12-31",
|
||||
@@ -394,17 +419,67 @@ class PublicStatusTests(unittest.TestCase):
|
||||
self.assertIn("G9", target_package["semantics"])
|
||||
|
||||
tushare = freshness["tushare-daily-market-data"]
|
||||
self.assertEqual(tushare["completedPartitions"], 35)
|
||||
self.assertEqual(tushare["plannedPartitions"], 428)
|
||||
self.assertAlmostEqual(tushare["completionRatio"], 35 / 428, places=6)
|
||||
self.assertEqual(tushare["marketDataThrough"], "1993-10-29")
|
||||
self.assertEqual(tushare["status"], "research_only")
|
||||
self.assertEqual(tushare["completedPartitions"], 96)
|
||||
self.assertEqual(tushare["plannedPartitions"], 96)
|
||||
self.assertEqual(tushare["completionRatio"], 1.0)
|
||||
self.assertEqual(tushare["completedRows"], 8_813_885)
|
||||
self.assertEqual(tushare["marketDataFrom"], "2018-01-02")
|
||||
self.assertEqual(tushare["marketDataThrough"], "2025-12-31")
|
||||
self.assertEqual(
|
||||
(tushare["verifiedSnapshotJobs"], tushare["snapshotJobs"]),
|
||||
(221, 221),
|
||||
)
|
||||
self.assertEqual(tushare["providerFiles"], 10_011)
|
||||
self.assertEqual(tushare["providerBytes"], 71_707_194)
|
||||
self.assertEqual(tushare["calendarSessions"], 1_942)
|
||||
self.assertEqual(tushare["instrumentCount"], 1_111)
|
||||
self.assertEqual(tushare["selectedDailyRows"], 1_975_455)
|
||||
self.assertEqual(tushare["observedMembershipSnapshots"], 97)
|
||||
self.assertEqual(tushare["effectiveMembershipSnapshots"], 96)
|
||||
self.assertFalse(tushare["sameSessionMembershipUse"])
|
||||
self.assertTrue(tushare["sourceStable"])
|
||||
self.assertTrue(tushare["stateDbChangedDuringVerification"])
|
||||
self.assertTrue(tushare["converterSourceMatchesCurrent"])
|
||||
self.assertEqual(
|
||||
tushare["providerDataVersion"],
|
||||
"5bf19d2da064357ad1802bca4bfa0c0505ed63fe2b24785ec6c56ecff1724963",
|
||||
)
|
||||
self.assertEqual(
|
||||
tushare["providerManifestSha256"],
|
||||
"27fb6fedb6a4b114eae2aba44505fb6c4d32c7a9ae81e58c724dfadb8dd10ed0",
|
||||
)
|
||||
self.assertEqual(tushare["sourceJobCount"], 221)
|
||||
self.assertEqual(
|
||||
tushare["lineageMatchedFields"],
|
||||
[
|
||||
"id",
|
||||
"path",
|
||||
"row_count",
|
||||
"byte_count",
|
||||
"sha256",
|
||||
"request_sha256",
|
||||
],
|
||||
)
|
||||
self.assertEqual(tushare["lineageMismatchCount"], 0)
|
||||
self.assertEqual(
|
||||
tushare["runSha256"],
|
||||
"803035a95ee9cf6f90af20cdf7e7024149b97e0113b585209a80f80048be87d3",
|
||||
)
|
||||
self.assertEqual(tushare["signalRows"], 176_615)
|
||||
self.assertEqual(tushare["portfolioDays"], 1_698)
|
||||
self.assertFalse(tushare["productionReady"])
|
||||
self.assertFalse(tushare["investmentValueClaim"])
|
||||
self.assertEqual(tushare["gateCredit"], [])
|
||||
source = (PROJECT / "docs" / "TUSHARE_LOCAL_DATA.md").read_text(
|
||||
encoding="utf-8"
|
||||
)
|
||||
self.assertRegex(source, r"`trade_cal`.*2026")
|
||||
self.assertNotIn("2026", tushare["marketDataThrough"])
|
||||
self.assertIn("daily", tushare["semantics"])
|
||||
self.assertIn("221/221", tushare["semantics"])
|
||||
self.assertIn(
|
||||
"tushare-csi500-research-run",
|
||||
tushare["evidenceLinks"],
|
||||
)
|
||||
|
||||
self.assertEqual(freshness["jqdata-real-snapshot"]["value"], "N/A")
|
||||
self.assertEqual(freshness["jqdata-real-snapshot"]["status"], "blocked")
|
||||
|
||||
@@ -0,0 +1,105 @@
|
||||
import hashlib
|
||||
import json
|
||||
from pathlib import Path
|
||||
import sys
|
||||
import unittest
|
||||
|
||||
|
||||
ROOT = Path(__file__).resolve().parents[1]
|
||||
sys.path.insert(0, str(ROOT))
|
||||
|
||||
from tools.tushare_lineage import (
|
||||
TushareLineageError,
|
||||
_compare_source_jobs,
|
||||
_validate_snapshot,
|
||||
)
|
||||
|
||||
|
||||
def _canonical(value):
|
||||
return json.dumps(
|
||||
value,
|
||||
ensure_ascii=False,
|
||||
sort_keys=True,
|
||||
separators=(",", ":"),
|
||||
allow_nan=False,
|
||||
).encode()
|
||||
|
||||
|
||||
class TushareLineageTest(unittest.TestCase):
|
||||
def _snapshot(self):
|
||||
job = {
|
||||
"api_name": "daily",
|
||||
"byte_count": 123,
|
||||
"completed_at": "2026-07-31T00:00:00Z",
|
||||
"id": 7,
|
||||
"observed": {
|
||||
"row_count": 11,
|
||||
"sha256": "a" * 64,
|
||||
"size_bytes": 123,
|
||||
},
|
||||
"partition_key": "month=2024-01",
|
||||
"path": "data/parquet/daily/month=2024-01.parquet",
|
||||
"request_sha256": "b" * 64,
|
||||
"row_count": 11,
|
||||
"sha256": "a" * 64,
|
||||
"shadowed_completed_count": 0,
|
||||
"verified": True,
|
||||
}
|
||||
selection_identity = {
|
||||
key: value
|
||||
for key, value in job.items()
|
||||
if key not in {"observed", "verified"}
|
||||
}
|
||||
snapshot = {
|
||||
"format": "quant-os.tushare-scoped-snapshot/v2",
|
||||
"selected_jobs": [job],
|
||||
"selection_sha256": hashlib.sha256(
|
||||
_canonical({"selected_jobs": [selection_identity]})
|
||||
).hexdigest(),
|
||||
"verification": {
|
||||
"double_verified_jobs": 1,
|
||||
"selected_jobs_still_latest_after_verification": True,
|
||||
"verified_jobs": 1,
|
||||
},
|
||||
}
|
||||
snapshot["manifest_sha256"] = hashlib.sha256(
|
||||
_canonical(snapshot)
|
||||
).hexdigest()
|
||||
return snapshot
|
||||
|
||||
def test_exact_source_lineage_and_tamper_detection(self):
|
||||
snapshot = self._snapshot()
|
||||
jobs = _validate_snapshot(snapshot)
|
||||
provider_job = {
|
||||
"api_name": "daily",
|
||||
"byte_count": 123,
|
||||
"job_id": 7,
|
||||
"partition_key": "month=2024-01",
|
||||
"path": "data/parquet/daily/month=2024-01.parquet",
|
||||
"request_sha256": "b" * 64,
|
||||
"row_count": 11,
|
||||
"sha256": "a" * 64,
|
||||
}
|
||||
_compare_source_jobs(jobs, [provider_job])
|
||||
provider_job["sha256"] = "c" * 64
|
||||
with self.assertRaisesRegex(
|
||||
TushareLineageError,
|
||||
"source job fields differ",
|
||||
):
|
||||
_compare_source_jobs(jobs, [provider_job])
|
||||
|
||||
snapshot["selected_jobs"][0]["path"] = "../escape.parquet"
|
||||
snapshot_without_hash = dict(snapshot)
|
||||
snapshot_without_hash.pop("manifest_sha256")
|
||||
snapshot["manifest_sha256"] = hashlib.sha256(
|
||||
_canonical(snapshot_without_hash)
|
||||
).hexdigest()
|
||||
with self.assertRaisesRegex(
|
||||
TushareLineageError,
|
||||
"stay relative",
|
||||
):
|
||||
_validate_snapshot(snapshot)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,379 @@
|
||||
import hashlib
|
||||
import importlib.util
|
||||
import json
|
||||
from pathlib import Path
|
||||
import sqlite3
|
||||
import sys
|
||||
import tempfile
|
||||
import unittest
|
||||
|
||||
|
||||
ROOT = Path(__file__).resolve().parents[1]
|
||||
sys.path.insert(0, str(ROOT))
|
||||
|
||||
from tools.tushare_snapshot import (
|
||||
TushareSnapshotError,
|
||||
create_snapshot,
|
||||
write_manifest,
|
||||
)
|
||||
|
||||
|
||||
def _manifest_hash(manifest):
|
||||
payload = dict(manifest)
|
||||
payload.pop("manifest_sha256")
|
||||
encoded = json.dumps(
|
||||
payload,
|
||||
ensure_ascii=False,
|
||||
sort_keys=True,
|
||||
separators=(",", ":"),
|
||||
allow_nan=False,
|
||||
).encode()
|
||||
return hashlib.sha256(encoded).hexdigest()
|
||||
|
||||
|
||||
class TushareSnapshotTest(unittest.TestCase):
|
||||
def _new_mirror(self, root: Path):
|
||||
mirror = root / "mirror"
|
||||
(mirror / "data/parquet").mkdir(parents=True)
|
||||
connection = sqlite3.connect(mirror / "data/state.sqlite3")
|
||||
connection.execute(
|
||||
"""
|
||||
CREATE TABLE jobs (
|
||||
id INTEGER PRIMARY KEY,
|
||||
api_name TEXT NOT NULL,
|
||||
partition_key TEXT NOT NULL,
|
||||
status TEXT NOT NULL,
|
||||
row_count INTEGER,
|
||||
byte_count INTEGER,
|
||||
sha256 TEXT,
|
||||
file_path TEXT,
|
||||
completed_at TEXT,
|
||||
params_json TEXT,
|
||||
fields TEXT
|
||||
)
|
||||
"""
|
||||
)
|
||||
return mirror, connection
|
||||
|
||||
def _completed(
|
||||
self,
|
||||
mirror: Path,
|
||||
connection: sqlite3.Connection,
|
||||
*,
|
||||
job_id: int,
|
||||
api: str,
|
||||
partition: str,
|
||||
content: bytes = b"not-a-parquet",
|
||||
path_suffix: str = "",
|
||||
row_count: int = 1,
|
||||
):
|
||||
safe_partition = (
|
||||
partition.replace("/", "_").replace("=", "-") + path_suffix
|
||||
)
|
||||
path = mirror / "data/parquet" / api / f"{safe_partition}.parquet"
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
path.write_bytes(content)
|
||||
connection.execute(
|
||||
"""
|
||||
INSERT INTO jobs
|
||||
(id, api_name, partition_key, status, row_count, byte_count,
|
||||
sha256, file_path, completed_at, params_json, fields)
|
||||
VALUES (?, ?, ?, 'completed', ?, ?, ?, ?, ?, '[]', NULL)
|
||||
""",
|
||||
(
|
||||
job_id,
|
||||
api,
|
||||
partition,
|
||||
row_count,
|
||||
path.stat().st_size,
|
||||
hashlib.sha256(content).hexdigest(),
|
||||
str(path),
|
||||
f"2018-03-{min(job_id, 28):02d}T00:00:00Z",
|
||||
),
|
||||
)
|
||||
return path
|
||||
|
||||
def _core_fixture(self, root: Path):
|
||||
mirror, connection = self._new_mirror(root)
|
||||
self._completed(
|
||||
mirror,
|
||||
connection,
|
||||
job_id=1,
|
||||
api="daily",
|
||||
partition="month=2018-01",
|
||||
path_suffix="-old",
|
||||
)
|
||||
self._completed(
|
||||
mirror,
|
||||
connection,
|
||||
job_id=2,
|
||||
api="daily",
|
||||
partition="month=2018-01",
|
||||
content=b"newer",
|
||||
)
|
||||
self._completed(
|
||||
mirror,
|
||||
connection,
|
||||
job_id=3,
|
||||
api="daily",
|
||||
partition="month=2018-02",
|
||||
)
|
||||
self._completed(
|
||||
mirror,
|
||||
connection,
|
||||
job_id=4,
|
||||
api="daily",
|
||||
partition="month=2018-03",
|
||||
)
|
||||
self._completed(
|
||||
mirror,
|
||||
connection,
|
||||
job_id=5,
|
||||
api="adj_factor",
|
||||
partition="month=2018-01",
|
||||
)
|
||||
self._completed(
|
||||
mirror,
|
||||
connection,
|
||||
job_id=6,
|
||||
api="adj_factor",
|
||||
partition="month=2018-02",
|
||||
)
|
||||
self._completed(
|
||||
mirror,
|
||||
connection,
|
||||
job_id=7,
|
||||
api="trade_cal",
|
||||
partition="exchange=SSE/year=2018",
|
||||
)
|
||||
self._completed(
|
||||
mirror,
|
||||
connection,
|
||||
job_id=8,
|
||||
api="stock_basic",
|
||||
partition="list_status=L",
|
||||
)
|
||||
self._completed(
|
||||
mirror,
|
||||
connection,
|
||||
job_id=9,
|
||||
api="index_daily",
|
||||
partition="ts_code=000905_SH/year=2018",
|
||||
)
|
||||
self._completed(
|
||||
mirror,
|
||||
connection,
|
||||
job_id=10,
|
||||
api="index_weight",
|
||||
partition="index_code=000905_SH/year=2018",
|
||||
)
|
||||
self._completed(
|
||||
mirror,
|
||||
connection,
|
||||
job_id=20,
|
||||
api="index_weight",
|
||||
partition="index_code=000905_SH/year=2017",
|
||||
)
|
||||
connection.execute(
|
||||
"""
|
||||
INSERT INTO jobs
|
||||
(id, api_name, partition_key, status, row_count, byte_count,
|
||||
sha256, file_path, completed_at)
|
||||
VALUES
|
||||
(11, 'index_daily', 'ts_code=931643_CNI/year=2018', 'running',
|
||||
NULL, NULL, NULL, NULL, NULL)
|
||||
"""
|
||||
)
|
||||
connection.commit()
|
||||
return mirror, connection
|
||||
|
||||
def test_scoped_selection_shadows_old_completed_and_excludes_other_index(self):
|
||||
with tempfile.TemporaryDirectory() as temporary:
|
||||
root = Path(temporary)
|
||||
mirror, connection = self._core_fixture(root)
|
||||
connection.close()
|
||||
|
||||
manifest = create_snapshot(
|
||||
mirror,
|
||||
start="2018-01-15",
|
||||
end="2018-02-03",
|
||||
verify=False,
|
||||
)
|
||||
|
||||
self.assertEqual(manifest["summary"]["selected_jobs"], 9)
|
||||
self.assertEqual(
|
||||
manifest["summary"]["shadowed_completed_jobs"],
|
||||
1,
|
||||
)
|
||||
selected_ids = {job["id"] for job in manifest["selected_jobs"]}
|
||||
self.assertIn(2, selected_ids)
|
||||
self.assertNotIn(1, selected_ids)
|
||||
self.assertNotIn(4, selected_ids)
|
||||
self.assertIn(20, selected_ids)
|
||||
self.assertEqual(
|
||||
manifest["scope"]["index_weight_anchor_year"],
|
||||
2017,
|
||||
)
|
||||
self.assertEqual(
|
||||
manifest["excluded_live_jobs"][0]["reason"],
|
||||
"different_index_code",
|
||||
)
|
||||
self.assertEqual(
|
||||
manifest["verification"]["mode"],
|
||||
"not_performed",
|
||||
)
|
||||
self.assertTrue(
|
||||
manifest["verification"][
|
||||
"selected_jobs_still_latest_after_verification"
|
||||
]
|
||||
)
|
||||
self.assertRegex(
|
||||
manifest["selected_jobs"][0]["request_sha256"],
|
||||
r"^[0-9a-f]{64}$",
|
||||
)
|
||||
self.assertNotIn(
|
||||
"params_json",
|
||||
manifest["selected_jobs"][0],
|
||||
)
|
||||
self.assertEqual(
|
||||
manifest["manifest_sha256"],
|
||||
_manifest_hash(manifest),
|
||||
)
|
||||
serialized = json.dumps(manifest)
|
||||
self.assertNotIn(str(root), serialized)
|
||||
self.assertTrue(
|
||||
all(
|
||||
not Path(job["path"]).is_absolute()
|
||||
for job in manifest["selected_jobs"]
|
||||
)
|
||||
)
|
||||
|
||||
output = root / "snapshot.json"
|
||||
self.assertEqual(write_manifest(output, manifest), output.resolve())
|
||||
self.assertEqual(
|
||||
json.loads(output.read_text())["manifest_sha256"],
|
||||
manifest["manifest_sha256"],
|
||||
)
|
||||
|
||||
def test_in_scope_running_job_fails_closed(self):
|
||||
with tempfile.TemporaryDirectory() as temporary:
|
||||
root = Path(temporary)
|
||||
mirror, connection = self._core_fixture(root)
|
||||
connection.execute(
|
||||
"""
|
||||
INSERT INTO jobs
|
||||
(id, api_name, partition_key, status, row_count, byte_count,
|
||||
sha256, file_path, completed_at)
|
||||
VALUES
|
||||
(12, 'daily', 'month=2018-02', 'running',
|
||||
NULL, NULL, NULL, NULL, NULL)
|
||||
"""
|
||||
)
|
||||
connection.commit()
|
||||
connection.close()
|
||||
with self.assertRaisesRegex(
|
||||
TushareSnapshotError,
|
||||
"contains running/deferred",
|
||||
):
|
||||
create_snapshot(
|
||||
mirror,
|
||||
start="2018-01-15",
|
||||
end="2018-02-03",
|
||||
verify=False,
|
||||
)
|
||||
|
||||
def test_missing_partition_fails_closed(self):
|
||||
with tempfile.TemporaryDirectory() as temporary:
|
||||
root = Path(temporary)
|
||||
mirror, connection = self._new_mirror(root)
|
||||
self._completed(
|
||||
mirror,
|
||||
connection,
|
||||
job_id=1,
|
||||
api="daily",
|
||||
partition="month=2018-01",
|
||||
)
|
||||
connection.commit()
|
||||
connection.close()
|
||||
with self.assertRaisesRegex(
|
||||
TushareSnapshotError,
|
||||
"missing daily:month=2018-02",
|
||||
):
|
||||
create_snapshot(
|
||||
mirror,
|
||||
start="2018-01-01",
|
||||
end="2018-02-28",
|
||||
apis=("daily",),
|
||||
verify=False,
|
||||
)
|
||||
|
||||
@unittest.skipUnless(
|
||||
importlib.util.find_spec("pyarrow") is not None,
|
||||
"pyarrow is needed for Parquet metadata verification",
|
||||
)
|
||||
def test_default_verification_detects_file_tampering(self):
|
||||
import pyarrow as pa
|
||||
import pyarrow.parquet as pq
|
||||
|
||||
with tempfile.TemporaryDirectory() as temporary:
|
||||
root = Path(temporary)
|
||||
mirror, connection = self._new_mirror(root)
|
||||
table = pa.table(
|
||||
{
|
||||
"ts_code": ["600000.SH", "000001.SZ"],
|
||||
"trade_date": ["20180102", "20180102"],
|
||||
}
|
||||
)
|
||||
parquet_path = (
|
||||
mirror / "data/parquet/daily/month-2018-01.parquet"
|
||||
)
|
||||
parquet_path.parent.mkdir(parents=True)
|
||||
pq.write_table(table, parquet_path)
|
||||
content = parquet_path.read_bytes()
|
||||
connection.execute(
|
||||
"""
|
||||
INSERT INTO jobs
|
||||
(id, api_name, partition_key, status, row_count, byte_count,
|
||||
sha256, file_path, completed_at, params_json, fields)
|
||||
VALUES
|
||||
(1, 'daily', 'month=2018-01', 'completed', 2, ?, ?, ?,
|
||||
'2018-02-01T00:00:00Z', '[]',
|
||||
'ts_code,trade_date')
|
||||
""",
|
||||
(
|
||||
len(content),
|
||||
hashlib.sha256(content).hexdigest(),
|
||||
str(parquet_path),
|
||||
),
|
||||
)
|
||||
connection.commit()
|
||||
connection.close()
|
||||
|
||||
manifest = create_snapshot(
|
||||
mirror,
|
||||
start="2018-01-01",
|
||||
end="2018-01-31",
|
||||
apis=("daily",),
|
||||
)
|
||||
self.assertEqual(manifest["verification"]["verified_jobs"], 1)
|
||||
self.assertEqual(
|
||||
manifest["verification"]["double_verified_jobs"],
|
||||
1,
|
||||
)
|
||||
self.assertEqual(
|
||||
manifest["selected_jobs"][0]["observed"]["row_count"],
|
||||
2,
|
||||
)
|
||||
|
||||
parquet_path.write_bytes(content + b"tampered")
|
||||
with self.assertRaisesRegex(TushareSnapshotError, "size mismatch"):
|
||||
create_snapshot(
|
||||
mirror,
|
||||
start="2018-01-01",
|
||||
end="2018-01-31",
|
||||
apis=("daily",),
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user