136 lines
5 KiB
Python
136 lines
5 KiB
Python
"""Tests for joining slim price estimates back onto properties.parquet.
|
|
|
|
estimate.py emits (Postcode, coalesced address, estimate columns) and
|
|
join_estimates attaches them by that natural key. These tests pin the
|
|
properties that make the key safe: it maps estimates onto the right rows
|
|
regardless of order (a shuffled estimates frame is the worst case), it is
|
|
idempotent, and it refuses a partial/foreign estimates file rather than
|
|
silently nulling prices.
|
|
"""
|
|
|
|
from pathlib import Path
|
|
|
|
import polars as pl
|
|
import pytest
|
|
|
|
from pipeline.transform.join_price_estimates import join_estimates
|
|
from pipeline.transform.price_estimation.utils import (
|
|
ESTIMATE_COLUMNS,
|
|
JOIN_ADDRESS,
|
|
JOIN_KEYS,
|
|
join_address_expr,
|
|
)
|
|
|
|
N = 200
|
|
|
|
|
|
def _write_merged(path: Path) -> pl.DataFrame:
|
|
"""properties.parquet with the natural-key columns, a sentinel order column,
|
|
and no estimates. Half the rows are sale-addressed, half EPC-only, so the
|
|
coalesce in the key is exercised; every coalesced address is unique."""
|
|
df = pl.DataFrame(
|
|
{
|
|
"Postcode": [f"AA{i % 7} {i % 9}AA" for i in range(N)],
|
|
"Address per Property Register": [
|
|
f"reg-{i}" if i % 2 == 0 else None for i in range(N)
|
|
],
|
|
"Address per EPC": [f"epc-{i}" if i % 2 == 1 else None for i in range(N)],
|
|
"order": list(range(N)),
|
|
"junk": [f"x{i}" for i in range(N)],
|
|
}
|
|
)
|
|
df.write_parquet(path)
|
|
return df
|
|
|
|
|
|
def _write_estimates(path: Path, merged_path: Path, *, shuffle: bool = True) -> None:
|
|
"""Estimates keyed by the natural key, derived from the merged file the way
|
|
estimate.py does. Estimate = order * 1000 so each row is checkable. Shuffled
|
|
by default to prove order-independence."""
|
|
est = (
|
|
pl.read_parquet(merged_path)
|
|
.with_columns(join_address_expr())
|
|
.with_columns(
|
|
(pl.col("order") * 1000).cast(pl.Float64).alias("Estimated current price"),
|
|
(pl.col("order") * 10).cast(pl.Int32).alias("Est. price per sqm"),
|
|
)
|
|
.select(*JOIN_KEYS, *ESTIMATE_COLUMNS)
|
|
)
|
|
if shuffle:
|
|
est = est.sample(fraction=1.0, shuffle=True, seed=7)
|
|
est.write_parquet(path)
|
|
|
|
|
|
def test_join_attaches_estimates_to_the_right_rows(tmp_path: Path):
|
|
props = tmp_path / "properties.parquet"
|
|
estimates = tmp_path / "price_estimates.parquet"
|
|
_write_merged(props)
|
|
_write_estimates(estimates, props)
|
|
|
|
written = join_estimates(props, estimates)
|
|
out = pl.read_parquet(props)
|
|
|
|
assert written == N
|
|
assert out.height == N
|
|
# Order preserved and the address-half of the key is not left behind.
|
|
assert out["order"].to_list() == list(range(N))
|
|
assert out["junk"].to_list() == [f"x{i}" for i in range(N)]
|
|
assert JOIN_ADDRESS not in out.columns
|
|
# Every row carries its own estimate, matched by key despite the shuffle.
|
|
assert out["Estimated current price"].to_list() == [float(i * 1000) for i in range(N)]
|
|
assert out["Est. price per sqm"].to_list() == [i * 10 for i in range(N)]
|
|
assert out["Estimated current price"].null_count() == 0
|
|
|
|
|
|
def test_rerun_is_idempotent(tmp_path: Path):
|
|
props = tmp_path / "properties.parquet"
|
|
estimates = tmp_path / "price_estimates.parquet"
|
|
_write_merged(props)
|
|
_write_estimates(estimates, props)
|
|
|
|
join_estimates(props, estimates)
|
|
first = pl.read_parquet(props)
|
|
join_estimates(props, estimates) # second run on the augmented file
|
|
second = pl.read_parquet(props)
|
|
|
|
assert second.equals(first)
|
|
assert second.columns.count("Estimated current price") == 1
|
|
assert second.columns.count("Est. price per sqm") == 1
|
|
|
|
|
|
def test_missing_estimate_is_rejected(tmp_path: Path):
|
|
"""A property with no matching estimate (diverged dwelling universe) must
|
|
fail loudly rather than silently leave its price null."""
|
|
props = tmp_path / "properties.parquet"
|
|
estimates = tmp_path / "price_estimates.parquet"
|
|
_write_merged(props)
|
|
_write_estimates(estimates, props)
|
|
# Drop one estimate so a property key is no longer covered.
|
|
pl.read_parquet(estimates).head(N - 1).write_parquet(estimates)
|
|
|
|
with pytest.raises(ValueError, match="no matching estimate"):
|
|
join_estimates(props, estimates)
|
|
|
|
|
|
def test_duplicate_key_is_rejected(tmp_path: Path):
|
|
props = tmp_path / "properties.parquet"
|
|
estimates = tmp_path / "price_estimates.parquet"
|
|
_write_merged(props)
|
|
_write_estimates(estimates, props)
|
|
# Force row 1's key to collide with row 0's.
|
|
est = pl.read_parquet(estimates).sort("Estimated current price")
|
|
row0 = est.row(0, named=True)
|
|
est = est.with_columns(
|
|
pl.when(pl.int_range(pl.len()) == 1)
|
|
.then(pl.lit(row0["Postcode"]))
|
|
.otherwise(pl.col("Postcode"))
|
|
.alias("Postcode"),
|
|
pl.when(pl.int_range(pl.len()) == 1)
|
|
.then(pl.lit(row0[JOIN_ADDRESS]))
|
|
.otherwise(pl.col(JOIN_ADDRESS))
|
|
.alias(JOIN_ADDRESS),
|
|
)
|
|
est.write_parquet(estimates)
|
|
|
|
with pytest.raises(ValueError, match="not unique"):
|
|
join_estimates(props, estimates)
|