wehub-resource-sync
216 lines
7.4 KiB
Python
216 lines
7.4 KiB
Python
"""Tests for utils/data_quality.py.
|
|
|
|
Pins:
|
|
- apply_max_symbols: seed determinism + edge cases (no-op when max<=0 or >=N).
|
|
Called by every loader; non-determinism would break reproducibility of tests
|
|
and notebooks that depend on a sampled subset.
|
|
- check_ohlc_invariants: correct detection of OHLC violations, graceful
|
|
handling of null values (TAQ no-trade bars), and missing-column tolerance.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import polars as pl
|
|
|
|
from utils.data_quality import (
|
|
apply_max_symbols,
|
|
check_ohlc_invariants,
|
|
describe_coverage,
|
|
null_rate,
|
|
)
|
|
|
|
|
|
def _make_prices(symbols: list[str], n_rows: int = 3) -> pl.DataFrame:
|
|
rows = []
|
|
for s in symbols:
|
|
for i in range(n_rows):
|
|
rows.append(
|
|
{
|
|
"symbol": s,
|
|
"timestamp": f"2024-01-{i + 1:02d}",
|
|
"open": 100.0 + i,
|
|
"high": 105.0 + i,
|
|
"low": 95.0 + i,
|
|
"close": 102.0 + i,
|
|
"volume": 1_000 * (i + 1),
|
|
}
|
|
)
|
|
return pl.DataFrame(rows)
|
|
|
|
|
|
# -----------------------------------------------------------------------------
|
|
# apply_max_symbols
|
|
# -----------------------------------------------------------------------------
|
|
|
|
|
|
def test_apply_max_symbols_zero_is_passthrough() -> None:
|
|
df = _make_prices(["A", "B", "C"])
|
|
out = apply_max_symbols(df, 0)
|
|
assert out.equals(df)
|
|
|
|
|
|
def test_apply_max_symbols_negative_is_passthrough() -> None:
|
|
df = _make_prices(["A", "B", "C"])
|
|
out = apply_max_symbols(df, -1)
|
|
assert out.equals(df)
|
|
|
|
|
|
def test_apply_max_symbols_exceeds_universe_is_passthrough() -> None:
|
|
df = _make_prices(["A", "B"])
|
|
out = apply_max_symbols(df, 10)
|
|
assert out.equals(df)
|
|
|
|
|
|
def test_apply_max_symbols_samples_requested_count() -> None:
|
|
df = _make_prices(["A", "B", "C", "D", "E"])
|
|
out = apply_max_symbols(df, 3)
|
|
assert out["symbol"].n_unique() == 3
|
|
|
|
|
|
def test_apply_max_symbols_is_seed_deterministic() -> None:
|
|
"""Same seed → same subset; critical for reproducible tests."""
|
|
df = _make_prices(["A", "B", "C", "D", "E", "F", "G", "H"])
|
|
|
|
first = apply_max_symbols(df, 3, seed=42)["symbol"].unique().sort().to_list()
|
|
second = apply_max_symbols(df, 3, seed=42)["symbol"].unique().sort().to_list()
|
|
assert first == second
|
|
|
|
|
|
def test_apply_max_symbols_different_seed_yields_different_sample() -> None:
|
|
df = _make_prices(["A", "B", "C", "D", "E", "F", "G", "H"])
|
|
|
|
s42 = set(apply_max_symbols(df, 3, seed=42)["symbol"].unique().to_list())
|
|
s7 = set(apply_max_symbols(df, 3, seed=7)["symbol"].unique().to_list())
|
|
# At least one sample differs — very high probability for k=3, n=8
|
|
assert s42 != s7
|
|
|
|
|
|
def test_apply_max_symbols_preserves_all_rows_per_symbol() -> None:
|
|
df = _make_prices(["A", "B", "C", "D"], n_rows=5)
|
|
out = apply_max_symbols(df, 2)
|
|
# Each selected symbol should keep all its rows (function filters by symbol set)
|
|
per_symbol = out.group_by("symbol").len()
|
|
assert per_symbol["len"].to_list() == [5, 5]
|
|
|
|
|
|
def test_apply_max_symbols_sort_then_sample_is_order_invariant() -> None:
|
|
"""Shuffling the input before sampling must yield the same subset: the
|
|
function sorts symbols before seeding the RNG so unstable loader order
|
|
(e.g., parquet partition order) can't perturb the selection."""
|
|
df_asc = _make_prices(["A", "B", "C", "D", "E"])
|
|
df_desc = df_asc.sort("symbol", descending=True)
|
|
|
|
s1 = apply_max_symbols(df_asc, 2, seed=42)["symbol"].unique().sort().to_list()
|
|
s2 = apply_max_symbols(df_desc, 2, seed=42)["symbol"].unique().sort().to_list()
|
|
assert s1 == s2
|
|
|
|
|
|
def test_apply_max_symbols_custom_symbol_col() -> None:
|
|
df = pl.DataFrame({"product": ["X", "Y", "Z"], "value": [1, 2, 3]})
|
|
out = apply_max_symbols(df, 2, symbol_col="product")
|
|
assert out["product"].n_unique() == 2
|
|
|
|
|
|
def test_apply_max_symbols_with_lazyframe_returns_lazyframe() -> None:
|
|
df = _make_prices(["A", "B", "C", "D"])
|
|
lf = df.lazy()
|
|
out = apply_max_symbols(lf, 2)
|
|
assert isinstance(out, pl.LazyFrame)
|
|
assert out.collect()["symbol"].n_unique() == 2
|
|
|
|
|
|
# -----------------------------------------------------------------------------
|
|
# check_ohlc_invariants
|
|
# -----------------------------------------------------------------------------
|
|
|
|
|
|
def test_check_ohlc_invariants_clean_data_is_100_percent() -> None:
|
|
df = _make_prices(["A"], n_rows=5)
|
|
invariants = check_ohlc_invariants(df)
|
|
# 6 checks expected: high_gte_low/open/close, low_lte_open/close, volume_non_negative
|
|
assert invariants.height == 6
|
|
assert (invariants["valid_pct"] == 100.0).all()
|
|
|
|
|
|
def test_check_ohlc_invariants_detects_high_below_low() -> None:
|
|
df = pl.DataFrame(
|
|
{
|
|
"open": [100.0, 100.0],
|
|
"high": [99.0, 105.0], # first row violates high >= low
|
|
"low": [100.0, 95.0],
|
|
"close": [99.5, 102.0],
|
|
"volume": [1_000, 1_000],
|
|
}
|
|
)
|
|
invariants = check_ohlc_invariants(df)
|
|
row = invariants.filter(pl.col("check") == "high_gte_low").row(0, named=True)
|
|
assert row["valid_pct"] == 50.0
|
|
assert row["applicable_rows"] == 2
|
|
|
|
|
|
def test_check_ohlc_invariants_ignores_null_rows() -> None:
|
|
"""Rows where any required col is null are excluded from the percentage."""
|
|
df = pl.DataFrame(
|
|
{
|
|
"open": [100.0, None, 100.0],
|
|
"high": [105.0, None, 102.0],
|
|
"low": [95.0, None, 95.0],
|
|
"close": [101.0, None, 101.0],
|
|
"volume": [1_000, 1_000, 1_000],
|
|
}
|
|
)
|
|
invariants = check_ohlc_invariants(df)
|
|
row = invariants.filter(pl.col("check") == "high_gte_low").row(0, named=True)
|
|
assert row["applicable_rows"] == 2 # 3 rows, 1 excluded for nulls
|
|
assert row["valid_pct"] == 100.0
|
|
|
|
|
|
def test_check_ohlc_invariants_detects_negative_volume() -> None:
|
|
df = pl.DataFrame(
|
|
{
|
|
"open": [100.0],
|
|
"high": [105.0],
|
|
"low": [95.0],
|
|
"close": [101.0],
|
|
"volume": [-5],
|
|
}
|
|
)
|
|
invariants = check_ohlc_invariants(df)
|
|
row = invariants.filter(pl.col("check") == "volume_non_negative").row(0, named=True)
|
|
assert row["valid_pct"] == 0.0
|
|
|
|
|
|
def test_check_ohlc_invariants_omits_volume_when_missing() -> None:
|
|
df = pl.DataFrame({"open": [100.0], "high": [105.0], "low": [95.0], "close": [101.0]})
|
|
invariants = check_ohlc_invariants(df)
|
|
assert "volume_non_negative" not in invariants["check"].to_list()
|
|
assert invariants.height == 5 # 5 price checks, no volume check
|
|
|
|
|
|
def test_check_ohlc_invariants_empty_df_returns_empty_result() -> None:
|
|
df = pl.DataFrame({"x": []})
|
|
invariants = check_ohlc_invariants(df)
|
|
assert invariants.height == 0
|
|
|
|
|
|
# -----------------------------------------------------------------------------
|
|
# Smaller coverage helpers
|
|
# -----------------------------------------------------------------------------
|
|
|
|
|
|
def test_describe_coverage_basic_shape() -> None:
|
|
df = _make_prices(["A", "B"], n_rows=3)
|
|
cov = describe_coverage(df)
|
|
assert cov["rows"] == 6
|
|
assert cov["assets"] == 2
|
|
assert cov["unique_times"] == 3
|
|
|
|
|
|
def test_null_rate_reports_per_column() -> None:
|
|
df = pl.DataFrame({"a": [1, None, 3], "b": [None, None, 3]})
|
|
rates = null_rate(df)
|
|
by_col = dict(zip(rates["column"], rates["null_pct"], strict=True))
|
|
# 1/3 ≈ 33.33 for a, 2/3 ≈ 66.66 for b
|
|
assert round(by_col["a"], 2) == 33.33
|
|
assert round(by_col["b"], 2) == 66.67
|