tests/test_extension.py

"""Phase 4 acceptance: extend flow, dedup, quota fallback, manifest atomicity.

Auth and GPU inference are both mocked, so these run offline and deterministically.
"""

from __future__ import annotations

import sys
import threading
from dataclasses import dataclass
from pathlib import Path

import numpy as np
import pandas as pd
import pytest

sys.path.insert(0, str(Path(__file__).resolve().parent.parent))

from src import config, extension, runtime
from src.extension import ExtensionError, add_model, estimate, extend_coverage
from src.store import SignalStore


@dataclass
class FakeProfile:
    username: str


def ohlcv(n=900, start="2023-01-01", freq="D"):
    rng = np.random.default_rng(4)
    close = pd.Series(100 * np.exp(np.cumsum(rng.normal(0.0004, 0.02, n))))
    ts = pd.date_range(start, periods=n, freq=freq, tz="UTC")
    return pd.DataFrame({
        "ts": ts, "open": close * 0.999, "high": close * 1.02,
        "low": close * 0.98, "close": close, "volume": 1000.0, "source": "test",
    })


@pytest.fixture
def wired(tmp_path, monkeypatch):
    """A fresh offline store wired into runtime, with GPU inference mocked."""
    store = SignalStore(repo_id=None, local_root=tmp_path / "store", offline=True)
    store.write_prices("BTC-USD", "1d", ohlcv())
    monkeypatch.setattr(runtime, "_store", store, raising=False)
    monkeypatch.setattr(runtime, "get_store", lambda: store)
    monkeypatch.setattr(extension, "HAS_SPACES", True)
    runtime.cache_clear()

    calls = {"n": 0}

    def fake_inference(model_id, family, values, ctx_len):
        calls["n"] += 1
        arr = np.asarray(values, dtype="float64")
        last = arr[:, -1]
        return {
            "q10": (last * 0.97).tolist(), "q50": last.tolist(),
            "q90": (last * 1.03).tolist(), "context_len": int(ctx_len),
            "revision": "deadbeef", "inference_version": "1.0.0+test.deadbeef",
        }

    monkeypatch.setattr(extension, "run_inference", fake_inference)
    return store, calls


# --------------------------------------------------------------------------
# Auth gating
# --------------------------------------------------------------------------


def test_anonymous_users_cannot_extend(wired):
    html, _ = extension.extend_ui("chronos-bolt-small", "BTC-USD", "1d",
                                  "2024-01-01", "2024-06-01", profile=None)
    assert "Sign in with Hugging Face" in html
    assert "spends your own GPU quota" in html


def test_anonymous_users_cannot_add_models(wired):
    html, _ = extension.add_model_ui("chronos", "amazon/chronos-bolt-small", profile=None)
    assert "Sign in" in html


def test_signed_in_user_can_extend(wired):
    store, calls = wired
    html, cov = extension.extend_ui("chronos-bolt-small", "BTC-USD", "1d",
                                    "2024-01-01", "2024-06-01",
                                    profile=FakeProfile("alice"))
    assert "Coverage extended by <b>@alice</b>" in html
    assert calls["n"] == 1
    assert not cov.empty


def test_contribution_is_attributed_to_the_user(wired):
    store, _ = wired
    extend_coverage("chronos-bolt-small", "BTC-USD", "1d",
                    "2024-01-01", "2024-06-01", username="bob")
    entries = store.load_manifest().find_signals(model_slug="chronos-bolt-small")
    assert entries and entries[0].contributed_by == "bob"


# --------------------------------------------------------------------------
# Dedup — never recompute covered ranges
# --------------------------------------------------------------------------


def test_second_identical_request_recomputes_nothing(wired):
    store, calls = wired
    first = extend_coverage("chronos-bolt-small", "BTC-USD", "1d",
                            "2024-01-01", "2024-06-01", username="alice")
    assert "Coverage extended" in first
    assert calls["n"] == 1

    second = extend_coverage("chronos-bolt-small", "BTC-USD", "1d",
                             "2024-01-01", "2024-06-01", username="alice")
    assert "already covered" in second
    assert calls["n"] == 1, "inference ran again for an already-covered range"


def test_estimate_reports_already_covered(wired):
    extend_coverage("chronos-bolt-small", "BTC-USD", "1d",
                    "2024-01-01", "2024-06-01", username="alice")
    est = estimate("chronos-bolt-small", "BTC-USD", "1d", "2024-02-01", "2024-05-01")
    assert est.already_covered


def test_estimate_counts_steps_for_an_uncovered_range(wired):
    est = estimate("chronos-bolt-small", "BTC-USD", "1d", "2024-01-01", "2024-06-01")
    assert est.steps > 0 and not est.already_covered


# --------------------------------------------------------------------------
# Guardrails
# --------------------------------------------------------------------------


def test_range_is_capped_per_timeframe(wired):
    est = estimate("chronos-bolt-small", "BTC-USD", "1d", "2015-01-01", "2025-01-01")
    assert est.capped
    assert (est.end - est.start).days <= config.CAPS.max_days["1d"]
    assert "cap" in est.note


def test_steps_never_exceed_the_per_run_ceiling(wired):
    est = estimate("chronos-bolt-small", "BTC-USD", "1d", "2023-01-01", "2025-01-01")
    assert est.steps <= config.CAPS.max_steps_per_run


@pytest.mark.parametrize("bad", [
    ("nope-model", "BTC-USD", "1d"),
    ("chronos-bolt-small", "DOGE-USD", "1d"),
    ("chronos-bolt-small", "BTC-USD", "3y"),
])
def test_unknown_selections_are_rejected(wired, bad):
    with pytest.raises(ExtensionError):
        estimate(bad[0], bad[1], bad[2], "2024-01-01", "2024-06-01")


def test_inverted_range_is_rejected(wired):
    with pytest.raises(ExtensionError, match="start must be before end"):
        estimate("chronos-bolt-small", "BTC-USD", "1d", "2024-06-01", "2024-01-01")


def test_extension_requires_existing_price_coverage(wired):
    with pytest.raises(ExtensionError, match="No cached prices"):
        estimate("chronos-bolt-small", "ETH-USD", "1d", "2024-01-01", "2024-06-01")


def test_add_model_rejects_hostile_ids(wired):
    for bad in ("../../etc/passwd", "no-slash", "owner/name;rm -rf /"):
        html = add_model("chronos", bad, username="alice")
        assert "not a valid Hub model id" in html or "not allowed" in html


def test_add_model_rejects_families_off_the_allow_list(wired):
    html = add_model("evil", "someone/backdoor", username="alice")
    assert "not on the allow-list" in html


def test_add_model_smoke_test_registers_on_success(wired):
    store, calls = wired
    html = add_model("chronos", "amazon/chronos-bolt-small", username="carol")
    assert "Smoke test passed" in html
    assert calls["n"] == 1
    slugs = {e.model_slug for e in store.load_manifest().signals.values()}
    assert "chronos-bolt-small" in slugs


# --------------------------------------------------------------------------
# Quota exhaustion
# --------------------------------------------------------------------------


def test_quota_exhaustion_renders_the_duplicate_space_fallback(wired, monkeypatch):
    def boom(*a, **k):
        raise RuntimeError("ZeroGPU quota exceeded for this account")

    monkeypatch.setattr(extension, "run_inference", boom)
    html = extend_coverage("chronos-bolt-small", "BTC-USD", "1d",
                           "2024-01-01", "2024-06-01", username="alice")
    assert "quota is exhausted" in html
    assert "duplicate=true" in html


@pytest.mark.parametrize("message", [
    "GPU task aborted", "ZeroGPU quota exceeded", "No GPU available right now",
])
def test_quota_errors_are_recognised(message):
    assert extension._is_quota_error(RuntimeError(message))


def test_ordinary_errors_are_not_mistaken_for_quota_errors(wired, monkeypatch):
    monkeypatch.setattr(extension, "run_inference",
                        lambda *a, **k: (_ for _ in ()).throw(ValueError("bad tensor shape")))
    html = extend_coverage("chronos-bolt-small", "BTC-USD", "1d",
                           "2024-01-01", "2024-06-01", username="alice")
    assert "quota" not in html.lower()
    assert "bad tensor shape" in html


# --------------------------------------------------------------------------
# Manifest atomicity under concurrent writers
# --------------------------------------------------------------------------


def test_parallel_extensions_do_not_corrupt_the_manifest(wired):
    """Simulated parallel writers must leave a valid, complete manifest."""
    store, calls = wired
    store.write_prices("ETH-USD", "1d", ohlcv())
    store.write_prices("SOL-USD", "1d", ohlcv())

    targets = [("BTC-USD", "alice"), ("ETH-USD", "bob"), ("SOL-USD", "carol")]
    errors: list[Exception] = []
    barrier = threading.Barrier(len(targets))

    def worker(asset, user):
        try:
            barrier.wait(timeout=10)  # maximise overlap
            extend_coverage("chronos-bolt-small", asset, "1d",
                            "2024-01-01", "2024-06-01", username=user)
        except Exception as e:
            errors.append(e)

    threads = [threading.Thread(target=worker, args=t) for t in targets]
    for t in threads:
        t.start()
    for t in threads:
        t.join(timeout=60)

    assert not errors, f"concurrent extensions raised: {errors}"

    # Manifest must still validate and hold every contribution.
    m = store.load_manifest(force=True)
    m.validate()
    assets = {e.asset for e in m.signals.values()}
    assert assets == {"BTC-USD", "ETH-USD", "SOL-USD"}
    contributors = {e.contributed_by for e in m.signals.values()}
    assert contributors == {"alice", "bob", "carol"}

    # And every referenced slice must actually exist and be readable.
    for e in m.signals.values():
        df = store.get_signals(e.model_slug, e.asset, e.timeframe)
        assert not df.empty, f"manifest references an empty slice for {e.key}"


def test_repeated_parallel_requests_for_the_same_slice_run_inference_once(wired):
    store, calls = wired
    barrier = threading.Barrier(4)

    def worker():
        barrier.wait(timeout=10)
        extend_coverage("chronos-bolt-small", "BTC-USD", "1d",
                        "2024-01-01", "2024-06-01", username="alice")

    threads = [threading.Thread(target=worker) for _ in range(4)]
    for t in threads:
        t.start()
    for t in threads:
        t.join(timeout=60)

    m = store.load_manifest(force=True)
    m.validate()
    entries = [e for e in m.signals.values() if e.asset == "BTC-USD"]
    assert len(entries) == 1, "duplicate manifest entries for the same slice"


def test_write_lock_serialises_commits():
    assert isinstance(extension._WRITE_LOCK, type(threading.Lock()))


# --------------------------------------------------------------------------
# CPU-only degradation
# --------------------------------------------------------------------------


def test_cpu_only_space_disables_extension_with_an_explanation(wired, monkeypatch):
    monkeypatch.setattr(extension, "HAS_SPACES", False)
    html, _ = extension.extend_ui("chronos-bolt-small", "BTC-USD", "1d",
                                  "2024-01-01", "2024-06-01",
                                  profile=FakeProfile("alice"))
    assert "running on CPU" in html
    assert "works normally" in html


def test_status_html_states_which_mode_the_space_is_in(monkeypatch):
    monkeypatch.setattr(extension, "HAS_SPACES", True)
    assert "ZEROGPU AVAILABLE" in extension.status_html()
    monkeypatch.setattr(extension, "HAS_SPACES", False)
    assert "CPU" in extension.status_html()