test_properties.py

"""Property-based tests.

The rest of the suite is example-based: pick an input, assert an outcome. That
is the right shape for testing a specific message or a known edge, and it has
one weakness, which is that it only ever checks the cases somebody thought to
write down.

The failures this library exists for are exactly the ones nobody wrote down. A
column order that changes and raises nothing took a model from 99.87% to 26.01%
here; the reordering that did it was not on anyone's list. So these tests state
the *property* instead and let Hypothesis go looking for a counterexample:
every permutation rather than three of them, every row count rather than the
round numbers, every bin configuration rather than the default.

When one fails it shrinks the input first, so what arrives is the smallest
frame that breaks the claim rather than the random one that happened to.
"""

from __future__ import annotations

import tempfile
from pathlib import Path

import numpy as np
import pytest
from hypothesis import HealthCheck, assume, given, settings
from hypothesis import strategies as st
from hypothesis.extra import numpy as hnp

from featureguard import (
    Action,
    DriftConfig,
    DriftDetector,
    SanityChecker,
    SanityConfig,
    Schema,
    SchemaError,
)

# Generating data is not free, and a few of these fit a detector per example.
# Deadlines cause flakes on a loaded CI runner rather than finding bugs.
slow = settings(deadline=None, max_examples=60,
                suppress_health_check=[HealthCheck.too_slow])


def names_of(n: int) -> list[str]:
    return [f"f{i}" for i in range(n)]


@st.composite
def table(draw, min_cols: int = 2, max_cols: int = 12,
          min_rows: int = 1, max_rows: int = 40):
    """A finite float64 table, with its column names."""
    n_cols = draw(st.integers(min_cols, max_cols))
    n_rows = draw(st.integers(min_rows, max_rows))
    data = draw(hnp.arrays(
        np.float64, (n_rows, n_cols),
        elements=st.floats(-1e6, 1e6, allow_nan=False, allow_infinity=False),
    ))
    return data, names_of(n_cols)


# ===========================================================================
# Schema
#
# The property that matters: a reordering is *always* caught. Not usually.
# ===========================================================================


@given(table())
@slow
def test_a_matching_batch_is_never_rejected(payload):
    """No false positives, at any width or row count."""
    data, names = payload
    Schema.from_names(names).validate(data, names=names)


@given(table(), st.randoms())
@slow
def test_every_column_permutation_is_caught(payload, rng):
    """The 99.87 -> 26.01 failure, over all permutations rather than a few."""
    data, names = payload
    shuffled = list(names)
    rng.shuffle(shuffled)
    assume(shuffled != names)  # the identity is not a reordering

    order = [names.index(n) for n in shuffled]
    with pytest.raises(SchemaError):
        Schema.from_names(names).validate(data[:, order], names=shuffled)


@given(table(), st.randoms())
@slow
def test_align_inverts_any_permutation(payload, rng):
    """`align` is the repair, so it must undo exactly what it detects."""
    data, names = payload
    shuffled = list(names)
    rng.shuffle(shuffled)
    order = [names.index(n) for n in shuffled]

    schema = Schema.from_names(names)
    restored = schema.align(data[:, order], shuffled)

    assert np.array_equal(restored, data)
    schema.validate(restored, names=names)  # and the repair passes the guard


@given(table(), st.integers(-6, 6))
@slow
def test_any_width_mismatch_is_caught(payload, delta):
    data, names = payload
    assume(delta != 0)
    width = len(names) + delta
    assume(width >= 1)

    padded = np.zeros((data.shape[0], width), dtype=np.float64)
    with pytest.raises(SchemaError):
        Schema.from_names(names).validate(padded)


@given(table())
@slow
def test_a_schema_survives_a_save_and_load_round_trip(payload):
    data, names = payload
    schema = Schema.from_names(names)

    with tempfile.TemporaryDirectory() as directory:
        path = Path(directory) / "schema.json"
        schema.save(path)
        reloaded = Schema.load(path)

    assert reloaded.to_dict() == schema.to_dict()
    reloaded.validate(data, names=names)


# ===========================================================================
# Drift
#
# The bug worth guarding: PSI has a sampling-noise floor of about
# (bins - 1) / rows, so on too few rows a stationary feature scores above the
# warn threshold and the detector alarms on data that never moved. `min_rows`
# is derived from that rather than guessed, and this is the claim.
# ===========================================================================


@given(st.integers(2, 40), st.floats(0.02, 0.5))
def test_the_noise_floor_always_sits_below_the_warn_threshold(n_bins, psi_warn):
    """The invariant the min_rows derivation exists to guarantee."""
    config = DriftConfig(n_bins=n_bins, psi_warn=psi_warn,
                         psi_alarm=max(psi_warn, 0.5))
    assert config.noise_floor < config.psi_warn


@given(
    n_bins=st.integers(4, 20),
    n_features=st.integers(1, 4),
    seed=st.integers(0, 2**31 - 1),
)
@slow
def test_stationary_data_never_alarms(n_bins, n_features, seed):
    """No false alarm when nothing moved, at the detector's own minimum rows."""
    rng = np.random.default_rng(seed)
    names = names_of(n_features)
    config = DriftConfig(n_bins=n_bins)

    detector = DriftDetector(names, config).fit(
        rng.normal(size=(4000, n_features)))

    rows = config.effective_min_rows
    detector.update(rng.normal(size=(rows, n_features)))

    report = detector.report()
    assert not report.alarms, (
        f"alarmed on data drawn from the reference distribution: "
        f"{[str(d) for d in report.alarms]}"
    )


@given(
    n_features=st.integers(1, 4),
    shift=st.floats(4.0, 12.0),
    seed=st.integers(0, 2**31 - 1),
)
@slow
def test_a_large_shift_is_always_detected(n_features, shift, seed):
    """The other half: it must not be so quiet that real drift slips past."""
    rng = np.random.default_rng(seed)
    names = names_of(n_features)
    config = DriftConfig()

    detector = DriftDetector(names, config).fit(
        rng.normal(size=(4000, n_features)))
    detector.update(
        rng.normal(size=(config.effective_min_rows, n_features)) + shift)

    assert not detector.report().ok


@given(
    window=st.integers(2, 8),
    n_batches=st.integers(1, 50),
    n_features=st.integers(1, 4),
    seed=st.integers(0, 2**31 - 1),
)
@slow
def test_memory_is_bounded_however_long_the_stream_runs(
    window, n_batches, n_features, seed
):
    """The streaming claim: cost is set by the window, not by stream length."""
    rng = np.random.default_rng(seed)
    names = names_of(n_features)
    config = DriftConfig(window_batches=window)
    detector = DriftDetector(names, config).fit(rng.normal(size=(2000, n_features)))

    for _ in range(n_batches):
        detector.update(rng.normal(size=(64, n_features)))

    # However many batches went through, only `window` of them are retained.
    assert detector.batches_in_window <= window
    assert detector.rows_in_window <= 64 * window

    ceiling = 8 * n_features * config.n_bins * (window + 1)
    assert detector.memory_bytes() <= ceiling, (
        f"{n_batches} batches through a {window}-batch window used "
        f"{detector.memory_bytes()} bytes, over the {ceiling}-byte ceiling"
    )


# ===========================================================================
# Sanity
# ===========================================================================


@given(table(min_rows=25, max_rows=60))
@slow
def test_clean_data_produces_no_actionable_issue(payload):
    data, names = payload
    # A constant column is legitimately reported as dead, and random floats
    # occasionally produce one. That check has its own tests.
    assume(all(len(np.unique(data[:, i])) > 1 for i in range(data.shape[1])))

    checker = SanityChecker(names).fit(data)
    report = checker.check(data)

    assert [i for i in report.issues if i.actionable] == []


@given(table(min_rows=5, max_rows=30),
       st.integers(0, 10 ** 6),
       st.sampled_from([np.nan, np.inf, -np.inf]))
@slow
def test_any_injected_non_finite_value_is_found(payload, position, bad):
    data, names = payload
    # WARN rather than the default RAISE, so the finding arrives as a report
    # to inspect instead of an exception.
    checker = SanityChecker(
        names, SanityConfig(on_non_finite=Action.WARN, check_dead=False)
    ).fit(np.nan_to_num(data))

    corrupted = data.copy()
    row = position % corrupted.shape[0]
    col = position % corrupted.shape[1]
    corrupted[row, col] = bad

    with pytest.warns(UserWarning):
        report = checker.check(corrupted)

    assert names[col] in report.affected_columns
    assert any(i.kind == "non_finite" for i in report.issues)


@given(table(min_rows=5, max_rows=30), st.integers(0, 10 ** 6))
@slow
def test_repair_always_returns_finite_data(payload, position):
    """Whatever goes in, what comes out of a repair is usable."""
    data, names = payload
    checker = SanityChecker(
        names,
        SanityConfig(
            on_non_finite=Action.REPAIR,
            on_out_of_bounds=Action.REPAIR,
            # Random floats throw up constant columns often enough to bury the
            # output in warnings about something this test is not asking about.
            check_dead=False,
        ),
    ).fit(np.nan_to_num(data))

    corrupted = data.copy()
    corrupted[position % corrupted.shape[0], position % corrupted.shape[1]] = np.nan

    report = checker.check(corrupted, repair=True)
    assert report.repaired is None or np.all(np.isfinite(report.repaired))