188 lines
9.1 KiB
Python
188 lines
9.1 KiB
Python
# partition_reader 通用分区读取核测试（spec §20.11.3：核内零语义+失败分级）
|
|
# 2026-09-17 TDD RED -> GREEN
|
|
import json
|
|
import os
|
|
|
|
import pandas as pd
|
|
import pytest
|
|
|
|
import sanguo_data.partition_reader as pr
|
|
|
|
|
|
@pytest.fixture
|
|
def tree(tmp_path):
|
|
root = tmp_path / "corpus"
|
|
(root / "news_meta" / "dt=2026-09-16").mkdir(parents=True)
|
|
(root / "news_meta" / "dt=2026-09-17").mkdir(parents=True)
|
|
(root / "news_meta" / "not-a-partition").mkdir(parents=True)
|
|
pd.DataFrame({"art_code": ["a1", "a2"], "stock_code": ["600519", "000001"]}) \
|
|
.to_parquet(root / "news_meta" / "dt=2026-09-16" / "part-0.parquet", index=False)
|
|
pd.DataFrame({"art_code": ["a3"], "stock_code": ["600519"]}) \
|
|
.to_parquet(root / "news_meta" / "dt=2026-09-17" / "part-0.parquet", index=False)
|
|
pd.DataFrame({"art_code": ["a4"], "stock_code": ["300750"]}) \
|
|
.to_parquet(root / "news_meta" / "dt=2026-09-17" / "part-1.parquet", index=False)
|
|
return str(root)
|
|
|
|
|
|
class TestDebrisSkip:
|
|
"""10-05 掉电案: 0 字节/截断坏件=写盘中断遗物 → skip+ERROR 不裸崩."""
|
|
|
|
def test_zero_byte_part_skipped(self, tree, caplog):
|
|
d = os.path.join(tree, "news_meta", "dt=2026-09-17")
|
|
open(os.path.join(d, "part-2.parquet"), "wb").close() # 0 字节残骸
|
|
with caplog.at_level("ERROR", logger="partition_reader"):
|
|
df = pr.read_partitions(tree, "news_meta")
|
|
assert df["art_code"].tolist() == ["a1", "a2", "a3", "a4"]
|
|
assert "0 字节" in caplog.text
|
|
|
|
def test_corrupt_part_with_content_still_raises(self, tree):
|
|
"""坏件含内容 → 继续 fail loud(integrity_gate 单查隔离依赖异常
|
|
记 error=; 0 字节残骸才 skip——1ea143e 实测打穿教训)."""
|
|
d = os.path.join(tree, "news_meta", "dt=2026-09-17")
|
|
with open(os.path.join(d, "part-3.parquet"), "wb") as f:
|
|
f.write(b"PAR1garbage-truncated")
|
|
with pytest.raises(Exception):
|
|
pr.read_partitions(tree, "news_meta")
|
|
|
|
|
|
class TestRoots:
|
|
def test_env_override(self, monkeypatch, tmp_path):
|
|
monkeypatch.setenv("SANGUO_CORPUS_ROOT", str(tmp_path))
|
|
assert pr.corpus_root() == str(tmp_path)
|
|
|
|
def test_default_posix(self, monkeypatch):
|
|
monkeypatch.delenv("SANGUO_CORPUS_ROOT", raising=False)
|
|
monkeypatch.setattr(pr, "_IS_WIN", False)
|
|
assert pr.corpus_root() == "/volume1/stock/corpus"
|
|
|
|
def test_default_win(self, monkeypatch):
|
|
monkeypatch.delenv("SANGUO_CORPUS_ROOT", raising=False)
|
|
monkeypatch.setattr(pr, "_IS_WIN", True)
|
|
assert pr.corpus_root() == r"C:\sanguo_vnpy_v2\data\corpus"
|
|
|
|
|
|
class TestListDates:
|
|
def test_sorted_dt_values_ignores_non_partition(self, tree):
|
|
assert pr.list_dates(tree, "news_meta") == ["2026-09-16", "2026-09-17"]
|
|
|
|
def test_missing_domain_empty(self, tree):
|
|
assert pr.list_dates(tree, "nope") == []
|
|
|
|
|
|
class TestReadPartitions:
|
|
def test_concat_all_partitions_and_part_files(self, tree):
|
|
df = pr.read_partitions(tree, "news_meta")
|
|
assert list(df["art_code"]) == ["a1", "a2", "a3", "a4"]
|
|
|
|
def test_start_end_filter_inclusive(self, tree):
|
|
df = pr.read_partitions(tree, "news_meta", start="2026-09-17", end="2026-09-17")
|
|
assert list(df["art_code"]) == ["a3", "a4"]
|
|
df = pr.read_partitions(tree, "news_meta", start="2026-09-16", end="2026-09-16")
|
|
assert list(df["art_code"]) == ["a1", "a2"]
|
|
|
|
def test_duplicate_rows_survive_no_silent_dedup(self, tmp_path):
|
|
# 核内零去重（因子线契约）：同键行跨分区重复必须原样透出
|
|
root = tmp_path / "d"
|
|
for dt in ("dt=2026-09-16", "dt=2026-09-17"):
|
|
(root / "x" / dt).mkdir(parents=True)
|
|
pd.DataFrame({"SECURITY_CODE": ["000001"], "REPORT_DATE": ["2026-06-30"]}) \
|
|
.to_parquet(root / "x" / dt / "part-0.parquet", index=False)
|
|
df = pr.read_partitions(str(root), "x")
|
|
assert len(df) == 2
|
|
|
|
def test_missing_root_empty_df(self, tmp_path):
|
|
df = pr.read_partitions(str(tmp_path / "absent"), "news_meta")
|
|
assert isinstance(df, pd.DataFrame) and df.empty
|
|
|
|
def test_empty_domain_dir_empty_df(self, tmp_path):
|
|
(tmp_path / "corpus" / "news_meta").mkdir(parents=True)
|
|
df = pr.read_partitions(str(tmp_path / "corpus"), "news_meta")
|
|
assert df.empty
|
|
|
|
def test_schema_column_drift_fails_loud(self, tmp_path):
|
|
root = tmp_path / "d"
|
|
(root / "x" / "dt=2026-09-16").mkdir(parents=True)
|
|
(root / "x" / "dt=2026-09-17").mkdir(parents=True)
|
|
pd.DataFrame({"a": [1], "b": [2]}).to_parquet(root / "x" / "dt=2026-09-16" / "p.parquet", index=False)
|
|
pd.DataFrame({"a": [1], "c": [3]}).to_parquet(root / "x" / "dt=2026-09-17" / "p.parquet", index=False)
|
|
with pytest.raises(pr.PartitionSchemaError):
|
|
pr.read_partitions(str(root), "x")
|
|
|
|
def test_schema_text_vs_num_conflict_fails_loud(self, tmp_path):
|
|
root = tmp_path / "d"
|
|
(root / "x" / "dt=2026-09-16").mkdir(parents=True)
|
|
(root / "x" / "dt=2026-09-17").mkdir(parents=True)
|
|
pd.DataFrame({"a": [1]}).to_parquet(root / "x" / "dt=2026-09-16" / "p.parquet", index=False)
|
|
pd.DataFrame({"a": ["x"]}).to_parquet(root / "x" / "dt=2026-09-17" / "p.parquet", index=False)
|
|
with pytest.raises(pr.PartitionSchemaError):
|
|
pr.read_partitions(str(root), "x")
|
|
|
|
def test_schema_null_vs_double_compatible(self, tmp_path):
|
|
# EM 真实形态：全空列 arrow null(object) ↔ double 跨 part 文件漂移=合法
|
|
root = tmp_path / "d"
|
|
(root / "x" / "dt=2026-06-30").mkdir(parents=True)
|
|
pd.DataFrame({"v": [None], "k": ["a"]}).to_parquet(
|
|
root / "x" / "dt=2026-06-30" / "part-0.parquet", index=False)
|
|
pd.DataFrame({"v": [1.5], "k": ["b"]}).to_parquet(
|
|
root / "x" / "dt=2026-06-30" / "part-1.parquet", index=False)
|
|
df = pr.read_partitions(str(root), "x")
|
|
assert len(df) == 2
|
|
assert df["v"].dtype == "float64" and df["v"].isna().iloc[0]
|
|
|
|
def test_schema_int_float_compatible_upcast(self, tmp_path):
|
|
# 数值族内互通：int64+float64 合并升位（金额元值 <2^53 无精度损）
|
|
root = tmp_path / "d"
|
|
(root / "x" / "dt=2026-06-30").mkdir(parents=True)
|
|
(root / "x" / "dt=2026-09-30").mkdir(parents=True)
|
|
pd.DataFrame({"v": [3]}).to_parquet(root / "x" / "dt=2026-06-30" / "p.parquet", index=False)
|
|
pd.DataFrame({"v": [2.5]}).to_parquet(root / "x" / "dt=2026-09-30" / "p.parquet", index=False)
|
|
df = pr.read_partitions(str(root), "x")
|
|
assert df["v"].dtype == "float64" and sorted(df["v"]) == [2.5, 3.0]
|
|
|
|
def test_schema_null_bridging_fails_loud(self, tmp_path):
|
|
# 因子线 review 必修#1：首文件全空列把 num 与 text 桥接，三段必须 raise
|
|
root = tmp_path / "d"
|
|
(root / "x" / "dt=2026-06-30").mkdir(parents=True)
|
|
pd.DataFrame({"v": [None], "k": ["a"]}).to_parquet(
|
|
root / "x" / "dt=2026-06-30" / "part-0.parquet", index=False)
|
|
pd.DataFrame({"v": [1], "k": ["b"]}).to_parquet(
|
|
root / "x" / "dt=2026-06-30" / "part-1.parquet", index=False)
|
|
pd.DataFrame({"v": ["txt"], "k": ["c"]}).to_parquet(
|
|
root / "x" / "dt=2026-06-30" / "part-2.parquet", index=False)
|
|
with pytest.raises(pr.PartitionSchemaError):
|
|
pr.read_partitions(str(root), "x")
|
|
|
|
def test_schema_null_between_conflicts_still_fails(self, tmp_path):
|
|
# resolved 不被 null 文件重置：num→null→text 同样 raise
|
|
root = tmp_path / "d"
|
|
(root / "x" / "dt=2026-06-30").mkdir(parents=True)
|
|
pd.DataFrame({"v": [1], "k": ["a"]}).to_parquet(
|
|
root / "x" / "dt=2026-06-30" / "part-0.parquet", index=False)
|
|
pd.DataFrame({"v": [None], "k": ["b"]}).to_parquet(
|
|
root / "x" / "dt=2026-06-30" / "part-1.parquet", index=False)
|
|
pd.DataFrame({"v": ["txt"], "k": ["c"]}).to_parquet(
|
|
root / "x" / "dt=2026-06-30" / "part-2.parquet", index=False)
|
|
with pytest.raises(pr.PartitionSchemaError):
|
|
pr.read_partitions(str(root), "x")
|
|
|
|
def test_schema_column_order_drift_tolerated(self, tmp_path):
|
|
# 同列集不同列序：按列名对齐不误报（策略线 review 观察顺势处理）
|
|
root = tmp_path / "d"
|
|
(root / "x" / "dt=2026-06-30").mkdir(parents=True)
|
|
pd.DataFrame({"a": [1], "b": ["x"]}).to_parquet(
|
|
root / "x" / "dt=2026-06-30" / "p0.parquet", index=False)
|
|
pd.DataFrame({"b": ["y"], "a": [2]}).to_parquet(
|
|
root / "x" / "dt=2026-06-30" / "p1.parquet", index=False)
|
|
df = pr.read_partitions(str(root), "x")
|
|
assert len(df) == 2 and df["a"].tolist() == [1, 2]
|
|
|
|
|
|
class TestWatermark:
|
|
def test_read_watermark_dict(self, tree):
|
|
with open(os.path.join(tree, pr.WATERMARK_NAME), "w") as f:
|
|
json.dump({"generated_at": "2026-09-17T16:00:00", "domains": {"news_meta": "2026-09-17"}}, f)
|
|
assert pr.read_watermark(tree)["domains"]["news_meta"] == "2026-09-17"
|
|
|
|
def test_read_watermark_absent_none(self, tree):
|
|
assert pr.read_watermark(tree) is None
|