106 lines
3.0 KiB
Python
106 lines
3.0 KiB
Python
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()
|