tests/test_store.py

"""Phase 0 acceptance: manifest round-trip, idempotent writes, price validation."""

from __future__ import annotations

import json

import pandas as pd
import pytest

from src import config
from src.store import (
    CoverageEntry,
    Manifest,
    PriceCoverage,
    SchemaError,
    SignalStore,
    empty_manifest,
    signal_key,
    validate_price_frame,
    validate_signal_frame,
)


# --------------------------------------------------------------------------
# Fixtures
# --------------------------------------------------------------------------


@pytest.fixture
def store(tmp_path) -> SignalStore:
    return SignalStore(repo_id=None, local_root=tmp_path / "store", offline=True)


def make_signals(n=30, start="2024-01-01", freq="D", base=100.0) -> pd.DataFrame:
    ts = pd.date_range(start, periods=n, freq=freq, tz="UTC")
    q50 = pd.Series([base + i for i in range(n)], dtype="float64")
    return pd.DataFrame({
        "ts": ts,
        "q10": q50 * 0.97,
        "q50": q50,
        "q90": q50 * 1.03,
        "context_len": 512,
        "inference_version": config.INFERENCE_VERSION,
    })


def make_prices(n=30, start="2024-01-01", freq="D", base=100.0) -> pd.DataFrame:
    ts = pd.date_range(start, periods=n, freq=freq, tz="UTC")
    close = pd.Series([base + i for i in range(n)], dtype="float64")
    return pd.DataFrame({
        "ts": ts,
        "open": close * 0.99,
        "high": close * 1.02,
        "low": close * 0.98,
        "close": close,
        "volume": 1000.0,
        "source": "test",
    })


# --------------------------------------------------------------------------
# Manifest round-trip
# --------------------------------------------------------------------------


def test_empty_manifest_round_trips():
    m = empty_manifest()
    again = Manifest.from_json(m.to_json())
    assert again.schema_version == config.MANIFEST_SCHEMA_VERSION
    assert again.signals == {} and again.prices == {}


def test_manifest_round_trip_preserves_entries():
    m = empty_manifest()
    m.upsert_signal(CoverageEntry(
        model_slug="chronos-bolt-small", model_id="amazon/chronos-bolt-small",
        model_revision="abc123", asset="BTC-USD", timeframe="1d",
        start_ts="2022-01-01T00:00:00Z", end_ts="2024-12-31T00:00:00Z",
        rows=1096, inference_version="1.0.0", last_updated="2026-08-15T00:00:00Z",
        contributed_by="seed",
    ))
    m.upsert_price(PriceCoverage(
        asset="BTC-USD", timeframe="1d", start_ts="2022-01-01T00:00:00Z",
        end_ts="2024-12-31T00:00:00Z", rows=1096, sources=["binance"],
        last_updated="2026-08-15T00:00:00Z",
    ))

    again = Manifest.from_json(m.to_json())
    assert again.to_dict()["signals"] == m.to_dict()["signals"]
    assert again.to_dict()["prices"] == m.to_dict()["prices"]

    e = again.get_signal("chronos-bolt-small", "abc123", "BTC-USD", "1d")
    assert e is not None and e.rows == 1096
    assert e.key == signal_key("chronos-bolt-small", "abc123", "BTC-USD", "1d")


def test_manifest_rejects_missing_schema_version():
    with pytest.raises(SchemaError, match="schema_version"):
        Manifest.from_json(json.dumps({"signals": {}, "prices": {}}))


def test_manifest_rejects_future_schema_version():
    with pytest.raises(SchemaError, match="newer than this app"):
        Manifest.from_json(json.dumps({"schema_version": 999}))


def test_manifest_rejects_malformed_json():
    with pytest.raises(SchemaError, match="not valid JSON"):
        Manifest.from_json("{not json")


def test_manifest_rejects_inverted_range():
    m = empty_manifest()
    bad = CoverageEntry(
        model_slug="m", model_id="o/m", model_revision="r", asset="BTC-USD",
        timeframe="1d", start_ts="2024-12-31T00:00:00Z", end_ts="2022-01-01T00:00:00Z",
        rows=1, inference_version="1.0.0", last_updated="x", contributed_by="seed",
    )
    with pytest.raises(SchemaError, match="start_ts after end_ts"):
        m.upsert_signal(bad)


def test_revision_is_part_of_identity():
    """Two revisions of the same model are distinct coverage, never merged."""
    m = empty_manifest()
    for rev in ("rev-a", "rev-b"):
        m.upsert_signal(CoverageEntry(
            model_slug="chronos", model_id="amazon/chronos", model_revision=rev,
            asset="BTC-USD", timeframe="1d", start_ts="2024-01-01T00:00:00Z",
            end_ts="2024-02-01T00:00:00Z", rows=32, inference_version="1.0.0",
            last_updated="x", contributed_by="seed",
        ))
    assert len(m.signals) == 2


# --------------------------------------------------------------------------
# Store persistence + idempotency
# --------------------------------------------------------------------------


def test_store_manifest_persists_to_disk(store, tmp_path):
    store.write_signals("m1", "org/m1", "rev1", "BTC-USD", "1d", make_signals())
    path = tmp_path / "store" / config.MANIFEST_PATH
    assert path.exists()
    reloaded = SignalStore(repo_id=None, local_root=tmp_path / "store", offline=True)
    assert reloaded.load_manifest().get_signal("m1", "rev1", "BTC-USD", "1d").rows == 30


def test_signal_write_is_idempotent(store):
    df = make_signals(n=30)
    first = store.write_signals("m1", "org/m1", "rev1", "BTC-USD", "1d", df)
    second = store.write_signals("m1", "org/m1", "rev1", "BTC-USD", "1d", df)

    assert first.rows == second.rows == 30
    got = store.get_signals("m1", "BTC-USD", "1d")
    assert len(got) == 30
    assert not got.index.duplicated().any()


def test_rewriting_a_slice_does_not_change_stored_values(store):
    """Append-only: an existing ts keeps its original value."""
    store.write_signals("m1", "org/m1", "rev1", "BTC-USD", "1d", make_signals(base=100.0))
    before = store.get_signals("m1", "BTC-USD", "1d")["q50"].tolist()

    conflicting = make_signals(base=999.0)
    store.write_signals("m1", "org/m1", "rev1", "BTC-USD", "1d", conflicting)
    after = store.get_signals("m1", "BTC-USD", "1d")["q50"].tolist()

    assert before == after


def test_extending_coverage_widens_range_and_adds_rows(store):
    store.write_signals("m1", "org/m1", "rev1", "BTC-USD", "1d",
                        make_signals(n=30, start="2024-01-01"))
    entry = store.write_signals("m1", "org/m1", "rev1", "BTC-USD", "1d",
                                make_signals(n=30, start="2024-02-01"))
    assert entry.rows == 60
    assert entry.start_ts.startswith("2024-01-01")
    assert entry.end_ts.startswith("2024-03-01")
    assert len(store.get_signals("m1", "BTC-USD", "1d")) == 60


def test_write_spanning_year_boundary_splits_files(store, tmp_path):
    store.write_signals("m1", "org/m1", "rev1", "BTC-USD", "1d",
                        make_signals(n=60, start="2023-12-10"))
    root = tmp_path / "store" / "signals" / "m1" / "BTC-USD" / "1d"
    assert (root / "2023.parquet").exists()
    assert (root / "2024.parquet").exists()
    assert len(store.get_signals("m1", "BTC-USD", "1d")) == 60


def test_coverage_across_years_counts_untouched_years(store):
    store.write_signals("m1", "org/m1", "rev1", "BTC-USD", "1d",
                        make_signals(n=20, start="2023-01-01"))
    entry = store.write_signals("m1", "org/m1", "rev1", "BTC-USD", "1d",
                                make_signals(n=20, start="2024-06-01"))
    assert entry.rows == 40


def test_has_coverage_and_missing_ranges(store):
    assert not store.has_coverage("m1", "rev1", "BTC-USD", "1d")
    store.write_signals("m1", "org/m1", "rev1", "BTC-USD", "1d",
                        make_signals(n=30, start="2024-01-01"))

    assert store.has_coverage("m1", "rev1", "BTC-USD", "1d",
                              "2024-01-05", "2024-01-20")
    assert not store.has_coverage("m1", "rev1", "BTC-USD", "1d",
                                  "2023-01-01", "2024-01-20")

    # Fully covered -> nothing to recompute. This is the extension dedup gate.
    assert store.missing_ranges("m1", "rev1", "BTC-USD", "1d",
                                "2024-01-05", "2024-01-20") == []
    gaps = store.missing_ranges("m1", "rev1", "BTC-USD", "1d",
                                "2023-06-01", "2024-06-01")
    assert len(gaps) == 2


def test_placeholder_coverage_can_be_excluded(store):
    store.write_signals("m1", "org/m1", "rev1", "BTC-USD", "1d", make_signals(),
                        inference_version=config.PLACEHOLDER_VERSION)
    assert store.has_coverage("m1", "rev1", "BTC-USD", "1d")
    assert not store.has_coverage("m1", "rev1", "BTC-USD", "1d",
                                  allow_placeholder=False)


def test_get_signals_respects_window(store):
    store.write_signals("m1", "org/m1", "rev1", "BTC-USD", "1d",
                        make_signals(n=30, start="2024-01-01"))
    got = store.get_signals("m1", "BTC-USD", "1d", "2024-01-10", "2024-01-14")
    assert len(got) == 5
    assert str(got.index[0].date()) == "2024-01-10"


def test_missing_coverage_returns_empty_frame_not_error(store):
    got = store.get_signals("nope", "BTC-USD", "1d", "2024-01-01", "2024-02-01")
    assert got.empty


def test_pending_files_are_tracked_for_commit(store):
    store.write_signals("m1", "org/m1", "rev1", "BTC-USD", "1d", make_signals())
    pending = store.pending
    assert config.MANIFEST_PATH in pending
    assert any(p.startswith("signals/m1/BTC-USD/1d/") for p in pending)


# --------------------------------------------------------------------------
# Signal frame validation
# --------------------------------------------------------------------------


def test_signal_validation_rejects_crossed_quantiles():
    df = make_signals(n=10)
    df.loc[3, "q10"] = df.loc[3, "q90"] + 5  # q10 > q50 > q90
    with pytest.raises(SchemaError, match="crossed quantiles"):
        validate_signal_frame(df)


def test_signal_validation_rejects_duplicate_timestamps():
    df = make_signals(n=10)
    df.loc[5, "ts"] = df.loc[4, "ts"]
    with pytest.raises(SchemaError, match="duplicate timestamps"):
        validate_signal_frame(df)


def test_signal_validation_rejects_missing_columns():
    df = make_signals(n=10).drop(columns=["q90"])
    with pytest.raises(SchemaError, match="missing columns"):
        validate_signal_frame(df)


def test_signal_validation_normalises_naive_timestamps_to_utc():
    df = make_signals(n=5)
    df["ts"] = df["ts"].dt.tz_localize(None)
    out = validate_signal_frame(df)
    assert str(out["ts"].dt.tz) == "UTC"


# --------------------------------------------------------------------------
# Price validation — injected bad rows must be caught
# --------------------------------------------------------------------------


def test_price_validation_accepts_clean_frame():
    out, report = validate_price_frame(make_prices(), "1d")
    assert report.ok and report.gaps == 0 and len(out) == 30


def test_price_validation_catches_negative_price():
    df = make_prices()
    df.loc[7, "close"] = -50.0
    with pytest.raises(SchemaError, match="non-positive"):
        validate_price_frame(df, "1d")


def test_price_validation_catches_zero_price():
    df = make_prices()
    df.loc[2, "open"] = 0.0
    with pytest.raises(SchemaError, match="non-positive"):
        validate_price_frame(df, "1d")


def test_price_validation_catches_duplicate_timestamps():
    df = make_prices()
    df.loc[9, "ts"] = df.loc[8, "ts"]
    with pytest.raises(SchemaError, match="duplicate timestamps"):
        validate_price_frame(df, "1d")


def test_price_validation_catches_inconsistent_ohlc():
    df = make_prices()
    df.loc[4, "high"] = df.loc[4, "low"] - 1.0
    with pytest.raises(SchemaError, match="inconsistent OHLC"):
        validate_price_frame(df, "1d")


def test_price_validation_catches_negative_volume():
    df = make_prices()
    df.loc[11, "volume"] = -1.0
    with pytest.raises(SchemaError, match="negative"):
        validate_price_frame(df, "1d")


def test_price_validation_catches_nan_price():
    df = make_prices()
    df.loc[6, "close"] = float("nan")
    with pytest.raises(SchemaError, match="NaN"):
        validate_price_frame(df, "1d")


def test_price_validation_reports_gaps_without_failing():
    """A gap is a fact about coverage, not a validation failure."""
    df = make_prices(n=30).drop(index=[10, 11, 12]).reset_index(drop=True)
    out, report = validate_price_frame(df, "1d")
    assert report.ok
    assert report.gaps == 1
    assert len(report.gap_ranges) == 1
    assert len(out) == 27


def test_price_validation_non_strict_collects_problems(make_bad=None):
    df = make_prices()
    df.loc[7, "close"] = -50.0
    out, report = validate_price_frame(df, "1d", strict=False)
    assert not report.ok
    assert any("non-positive" in p for p in report.problems)


def test_price_write_and_read_round_trip(store):
    cov = store.write_prices("BTC-USD", "1d", make_prices(n=40))
    assert cov.rows == 40 and cov.sources == ["test"]
    got = store.get_prices("BTC-USD", "1d")
    assert len(got) == 40
    assert list(got.columns) == ["open", "high", "low", "close", "volume", "source"]


def test_price_write_is_idempotent(store):
    df = make_prices(n=40)
    store.write_prices("BTC-USD", "1d", df)
    cov = store.write_prices("BTC-USD", "1d", df)
    assert cov.rows == 40
    assert len(store.get_prices("BTC-USD", "1d")) == 40


# --------------------------------------------------------------------------
# Redundant / interfering data
#
# The parquet path is keyed on model slug but the manifest is keyed on slug AND
# revision, so two revisions share one file. Without superseding, the manifest
# would record the new revision while the file still held the old numbers.
# --------------------------------------------------------------------------


def test_rewriting_the_same_version_is_idempotent(store):
    """A repeated seed or extension must change nothing at all."""
    store.write_signals("m", "org/m", "revA", "BTC-USD", "1d",
                        make_signals(base=100.0), inference_version="1.0.0")
    before = store.get_signals("m", "BTC-USD", "1d")["q50"].tolist()
    store.write_signals("m", "org/m", "revA", "BTC-USD", "1d",
                        make_signals(base=100.0), inference_version="1.0.0")
    after = store.get_signals("m", "BTC-USD", "1d")
    assert after["q50"].tolist() == before
    assert len(after) == 30


def test_a_new_revision_supersedes_the_old_numbers(store):
    store.write_signals("m", "org/m", "revA", "BTC-USD", "1d",
                        make_signals(base=100.0), inference_version="1.0.0")
    store.write_signals("m", "org/m", "revB", "BTC-USD", "1d",
                        make_signals(base=999.0), inference_version="2.0.0")

    got = store.get_signals("m", "BTC-USD", "1d")
    assert got["q50"].iloc[0] == 999.0, "manifest would claim revB but hold revA"
    assert set(got["inference_version"]) == {"2.0.0"}
    assert not got.index.duplicated().any()


def test_real_output_supersedes_a_placeholder_slice(store):
    """A PLACEHOLDER slice must never shadow real inference forever."""
    store.write_signals("p", "org/p", "PLACEHOLDER", "BTC-USD", "1d",
                        make_signals(base=1.0),
                        inference_version=config.PLACEHOLDER_VERSION)
    store.write_signals("p", "org/p", "revReal", "BTC-USD", "1d",
                        make_signals(base=500.0), inference_version="1.0.0")

    got = store.get_signals("p", "BTC-USD", "1d")
    assert got["q50"].iloc[0] == 500.0
    assert config.PLACEHOLDER_VERSION not in set(got["inference_version"])


def test_the_written_version_is_authoritative_over_the_frame(store):
    """The kwarg and the frame column can disagree; the kwarg records the
    manifest entry, so it has to win or supersede compares the wrong value."""
    df = make_signals(base=7.0)
    df["inference_version"] = "stale-value-from-the-caller"
    store.write_signals("m", "org/m", "revA", "BTC-USD", "1d", df,
                        inference_version="1.0.0")

    got = store.get_signals("m", "BTC-USD", "1d")
    assert set(got["inference_version"]) == {"1.0.0"}


def test_superseding_leaves_untouched_timestamps_alone(store):
    """Only the overlapping instants are replaced, not the whole file."""
    store.write_signals("m", "org/m", "revA", "BTC-USD", "1d",
                        make_signals(n=30, start="2024-01-01", base=100.0),
                        inference_version="1.0.0")
    # A newer version covering only the first 10 bars.
    store.write_signals("m", "org/m", "revB", "BTC-USD", "1d",
                        make_signals(n=10, start="2024-01-01", base=999.0),
                        inference_version="2.0.0")

    got = store.get_signals("m", "BTC-USD", "1d")
    assert len(got) == 30, "non-overlapping rows must survive"
    assert got["q50"].iloc[0] == 999.0
    assert got["inference_version"].iloc[0] == "2.0.0"
    assert got["inference_version"].iloc[-1] == "1.0.0"


def test_price_rows_keep_first_source_on_collision(store):
    """Prices have no version to compare, so first-writer-wins stands. Two
    providers disagreeing about one bar is ambiguous, not a supersede."""
    a = make_prices(n=10)
    a["source"] = "binance"
    b = make_prices(n=10, base=500.0)
    b["source"] = "coinbase"
    store.write_prices("BTC-USD", "1d", a)
    store.write_prices("BTC-USD", "1d", b)

    got = store.get_prices("BTC-USD", "1d")
    assert len(got) == 10
    assert set(got["source"]) == {"binance"}