tests/test_acceptance.py

"""The six acceptance tests from the build spec.

Each test below is numbered to match the spec. They are kept in one file, on
purpose: these are the conditions the build was accepted against, and someone
changing the pipeline later should be able to read the whole contract in one
sitting rather than reconstructing it from six files.

Everything here runs offline. The classifier tests use a mock that never opens
a socket, so the guards are tested without depending on a live endpoint
answering the same way twice.
"""

from __future__ import annotations

import json

import pytest

from pipeline import classify, harvest, schema, store
from pipeline.schema import SchemaError
from src import atlas, lineage
from tests.conftest import make_row


# ==========================================================================
# 1. Harvest resumes correctly from checkpoint after simulated interruption
# ==========================================================================


class FlakyApi:
    """A fake Hub that raises partway through, then works on the second run.

    Stands in for the real interruption -- a killed process, a dropped
    connection -- without needing to actually kill anything.
    """

    def __init__(self, fail_after=None):
        self.fail_after = fail_after
        self.calls = []
        self.token = None

    def model_info(self, model_id, **kwargs):
        self.calls.append(model_id)
        if self.fail_after is not None and len(self.calls) > self.fail_after:
            raise KeyboardInterrupt("simulated interruption")
        return _FakeInfo(model_id)


class _FakeInfo:
    def __init__(self, model_id):
        self.id = model_id
        self.author = model_id.split("/")[0]
        self.tags = ["license:apache-2.0"]
        self.pipeline_tag = "text-classification"
        self.downloads = 500
        self.downloads_all_time = 5000
        self.likes = 5
        self.created_at = None
        self.last_modified = None
        self.library_name = "transformers"
        self.card_data = None
        self.gated = False
        self.safetensors = None
        self.siblings = [type("S", (), {"rfilename": "model.safetensors",
                                        "size": 10_000})()]


def test_1_harvest_resumes_from_checkpoint(tmp_path, monkeypatch):
    """Interrupt enrichment, restart, and only the unfinished models re-fetch."""
    monkeypatch.setattr(harvest, "fetch_readme", lambda *a, **k: "card text")
    monkeypatch.setattr(harvest, "CHECKPOINT_EVERY", 2)

    candidates = [f"org/model-{i}" for i in range(10)]
    path = tmp_path / "checkpoint.json"
    throttle = harvest.Throttle(min_interval=0)

    # --- first run: dies after 5 models -----------------------------------
    api = FlakyApi(fail_after=5)
    checkpoint = harvest.Checkpoint(path).load()
    with pytest.raises(KeyboardInterrupt):
        for _ in harvest.enrich(api, candidates, checkpoint, throttle):
            pass

    # The checkpoint flushes every 2 models, so at least 4 survived the crash.
    reloaded = harvest.Checkpoint(path).load()
    done_after_crash = len(reloaded.rows)
    assert 4 <= done_after_crash <= 5, done_after_crash
    assert path.exists()

    # --- second run: resumes -----------------------------------------------
    api2 = FlakyApi()
    for _ in harvest.enrich(api2, candidates, reloaded, throttle):
        pass

    # Every candidate is now enriched...
    assert len(reloaded.rows) == 10
    assert set(reloaded.rows) == set(candidates)

    # ...and the second run did not re-fetch anything the first one finished.
    refetched = set(api2.calls) & set(list(reloaded.rows)[:done_after_crash])
    already_done = set(harvest.Checkpoint(path).load().rows)
    assert len(api2.calls) == 10 - done_after_crash, (
        f"resumed run re-fetched {len(api2.calls)} models, "
        f"expected {10 - done_after_crash}")
    assert already_done == set(candidates)


def test_1b_corrupt_checkpoint_starts_fresh_instead_of_crashing(tmp_path):
    """A truncated checkpoint must not take the whole run down with it."""
    path = tmp_path / "checkpoint.json"
    path.write_text('{"rows": [{"id": "a/b"}], "candid')  # cut mid-write

    checkpoint = harvest.Checkpoint(path).load()
    assert checkpoint.rows == {}
    assert checkpoint.candidates == []


# ==========================================================================
# 2. Classifier schema-validation rejects and retries a malformed response
# ==========================================================================


class MockCompletions:
    """Returns queued responses in order, recording every prompt it received."""

    def __init__(self, responses):
        self.responses = list(responses)
        self.prompts = []

    def __call__(self, messages):
        self.prompts.append(messages)
        if not self.responses:
            raise AssertionError("mock ran out of responses")
        return self.responses.pop(0)


def make_classifier(responses, monkeypatch):
    """An InferenceClassifier whose transport is the mock, not the network."""
    obj = classify.InferenceClassifier.__new__(classify.InferenceClassifier)
    obj.model = "mock-model"
    obj.name = "mock"
    obj.max_tokens = 400
    obj.retry_sleep = 0
    obj.rate_limit_retries = 1
    obj.reset_usage()
    mock = MockCompletions(responses)
    obj._complete = mock
    return obj, mock


VALID = json.dumps({
    "relevant": "yes", "task": "sentiment", "asset_class": "equities",
    "training_data_summary": "Financial PhraseBank.", "maintained": True,
    "has_eval": True, "red_flags": [],
})


def test_2_malformed_response_is_rejected_then_retried(monkeypatch):
    """A bad first response is retried once, and the good retry is accepted."""
    clf, mock = make_classifier(["not json at all", VALID], monkeypatch)

    result = clf.classify(make_row("a/b"))

    assert len(mock.prompts) == 2, "expected exactly one retry"
    assert result.relevant == "yes"
    assert result.task == "sentiment"

    # The retry must quote the rejection back, so the model can correct itself.
    retry_text = mock.prompts[1][-1]["content"]
    assert "rejected" in retry_text.lower()


def test_2b_still_malformed_after_retry_becomes_unclear(monkeypatch):
    """Two bad responses become `unclear` -- never a crash, never a drop."""
    clf, mock = make_classifier(["garbage", "{still: not valid"], monkeypatch)

    result = clf.classify(make_row("a/b"))

    assert len(mock.prompts) == 2, "must not retry more than once"
    assert result.relevant == "unclear"
    assert result.task == "other"
    assert any("malformed" in f for f in result.red_flags)


def test_2c_out_of_vocabulary_values_are_rejected(monkeypatch):
    """An invented category is a validation failure, not a new category."""
    invented = json.dumps({
        "relevant": "yes", "task": "risk_scoring",     # not in the taxonomy
        "asset_class": "equities", "training_data_summary": "",
        "maintained": True, "has_eval": False, "red_flags": [],
    })
    clf, mock = make_classifier([invented, VALID], monkeypatch)

    result = clf.classify(make_row("a/b"))

    assert len(mock.prompts) == 2
    assert result.task == "sentiment"
    # The rejection names the offending field -- that is what makes the retry
    # useful rather than a coin flip.
    assert "task" in mock.prompts[1][-1]["content"]


def test_2d_transport_failure_becomes_unclear_not_an_exception():
    """A dead endpoint degrades to `unclear`; it does not take the run down."""
    obj = classify.InferenceClassifier.__new__(classify.InferenceClassifier)
    obj.model, obj.name = "mock", "mock"
    obj.max_tokens, obj.retry_sleep, obj.rate_limit_retries = 400, 0, 1
    obj.reset_usage()

    def boom(messages):
        raise ConnectionError("endpoint down")

    obj._complete = boom
    result = obj.classify(make_row("a/b"))

    assert result.relevant == "unclear"
    assert any("inference error" in f for f in result.red_flags)


@pytest.mark.parametrize("payload,reason", [
    ('{"relevant": "yes"}', "missing fields"),
    ('{"relevant": "maybe", "task": "sentiment", "asset_class": "general",'
     ' "training_data_summary": "", "maintained": true, "has_eval": true,'
     ' "red_flags": []}', "relevance not in vocabulary"),
    ('{"relevant": "yes", "task": "sentiment", "asset_class": "mars",'
     ' "training_data_summary": "", "maintained": true, "has_eval": true,'
     ' "red_flags": []}', "asset class not in vocabulary"),
    ('{"relevant": "yes", "task": "sentiment", "asset_class": "general",'
     ' "training_data_summary": "", "maintained": "sort of", "has_eval": true,'
     ' "red_flags": []}', "non-boolean"),
    ('[]', "not an object"),
    ('', "empty"),
])
def test_2e_validator_rejects_bad_payloads(payload, reason):
    with pytest.raises(SchemaError):
        schema.parse(payload)


def test_2f_validator_accepts_the_shapes_models_actually_emit():
    """Fenced JSON and a bare-string red flag are formatting, not refusal."""
    fenced = '```json\n' + VALID + '\n```'
    assert schema.parse(fenced).relevant == "yes"

    prose = "Here you go:\n" + VALID + "\nHope that helps!"
    assert schema.parse(prose).task == "sentiment"

    single_flag = json.loads(VALID)
    single_flag["red_flags"] = "no license"
    assert schema.validate(single_flag).red_flags == ["no license"]


# ==========================================================================
# 3. Known-answer classification
# ==========================================================================

FINBERT_CARD = """
# FinBERT

FinBERT is a pre-trained NLP model to analyze sentiment of financial text. It
is built by further training the BERT language model in the finance domain,
using a large financial corpus and thereby fine-tuning it for financial
sentiment classification. Financial PhraseBank by Malo et al. (2014) is used
for fine-tuning.
"""

# The decoy: "stock" meaning warehouse inventory. A classifier keying on the
# word rather than the meaning fails this, which is exactly what it is for.
DECOY_CARD = """
# Stock Level Predictor

Predicts warehouse stock levels for retail inventory management and restocking
schedules. Trained on internal warehouse SKU movement logs from three
distribution centres. Helps supply chain teams avoid stockouts of physical
goods on shelves. Nothing to do with equities, securities or financial markets.
"""


def test_3_known_answer_finbert_is_relevant_sentiment(monkeypatch):
    clf, _ = make_classifier([VALID], monkeypatch)
    row = make_row("ProsusAI/finbert", readme=FINBERT_CARD)

    result = clf.classify(row)

    assert result.relevant == "yes"
    assert result.task == "sentiment"


def test_3b_known_answer_decoy_is_not_relevant(monkeypatch):
    """A planted non-finance "stock" model must be classified `no`."""
    answer = json.dumps({
        "relevant": "no", "task": "other", "asset_class": "general",
        "training_data_summary": "Warehouse SKU movement logs.",
        "maintained": True, "has_eval": False, "red_flags": [],
    })
    clf, _ = make_classifier([answer], monkeypatch)
    row = make_row("warehouse-ai/stock-level-predictor", readme=DECOY_CARD)

    result = clf.classify(row)
    assert result.relevant == "no"


def test_3c_irrelevant_models_are_filtered_out_of_the_index():
    """`relevant: no` never reaches atlas.parquet; `unclear` does."""
    harvested = {
        "ProsusAI/finbert": {"id": "ProsusAI/finbert", "author": "ProsusAI"},
        "warehouse-ai/stock-level-predictor": {
            "id": "warehouse-ai/stock-level-predictor", "author": "warehouse-ai"},
        "mystery/unreadable": {"id": "mystery/unreadable", "author": "mystery"},
    }
    classifications = {
        "ProsusAI/finbert": {"relevant": "yes", "task": "sentiment",
                             "asset_class": "general", "training_data_summary": "",
                             "maintained": True, "has_eval": True, "red_flags": []},
        "warehouse-ai/stock-level-predictor": {
            "relevant": "no", "task": "other", "asset_class": "general",
            "training_data_summary": "", "maintained": True, "has_eval": False,
            "red_flags": []},
        "mystery/unreadable": {"relevant": "unclear", "task": "other",
                               "asset_class": "general", "training_data_summary": "",
                               "maintained": False, "has_eval": False,
                               "red_flags": []},
    }

    rows = store.build_rows(harvested, classifications)
    ids = {r["id"] for r in rows}

    assert "ProsusAI/finbert" in ids
    assert "warehouse-ai/stock-level-predictor" not in ids, "decoy leaked into the index"
    assert "mystery/unreadable" in ids, "unclear must be kept, not dropped"
    assert all(r["auto_classified"] for r in rows)


# ==========================================================================
# 4. Reclassification cache: unchanged models are not re-sent
# ==========================================================================


class CountingClassifier:
    name = "counting"

    def __init__(self):
        self.seen = []

    def classify(self, row):
        self.seen.append(row["id"])
        return schema.Classification(relevant="yes", task="sentiment",
                                     asset_class="general")


def test_4_unchanged_model_is_not_reclassified(tmp_path):
    """Second run over identical rows sends nothing to the model."""
    inner = CountingClassifier()
    cache_path = tmp_path / "cache.json"
    rows = [make_row("a/one"), make_row("b/two")]

    first = classify.CachedClassifier(inner, cache_path)
    for row in rows:
        first.classify(row)
    first.save()

    assert len(inner.seen) == 2
    assert first.misses == 2 and first.hits == 0

    # A fresh wrapper, reading the cache off disk -- as the weekly job does.
    second = classify.CachedClassifier(CountingClassifier(), cache_path)
    for row in rows:
        second.classify(row)

    assert second.inner.seen == [], "unchanged models were re-sent to the LLM"
    assert second.hits == 2 and second.misses == 0


def test_4b_changed_last_modified_forces_reclassification(tmp_path):
    """A new commit changes the key, so the model is read again."""
    cache_path = tmp_path / "cache.json"
    row = make_row("a/one", last_modified="2026-01-01T00:00:00+00:00")

    first = classify.CachedClassifier(CountingClassifier(), cache_path)
    first.classify(row)
    first.save()

    moved = dict(row, last_modified="2026-07-01T00:00:00+00:00")
    second = classify.CachedClassifier(CountingClassifier(), cache_path)
    second.classify(moved)

    assert second.inner.seen == ["a/one"]
    assert second.misses == 1


def test_4c_cache_key_is_id_plus_last_modified():
    a = classify.cache_key({"id": "x/y", "last_modified": "2026-01-01"})
    b = classify.cache_key({"id": "x/y", "last_modified": "2026-02-01"})
    c = classify.cache_key({"id": "x/z", "last_modified": "2026-01-01"})
    assert a != b and a != c


def test_4d_cache_entry_from_an_older_schema_is_discarded(tmp_path):
    """A cache written before a schema change must not poison the index."""
    cache_path = tmp_path / "cache.json"
    row = make_row("a/one")
    cache_path.write_text(json.dumps({
        classify.cache_key(row): {"relevant": "yes", "task": "risk_scoring"},
    }))

    wrapper = classify.CachedClassifier(CountingClassifier(), cache_path)
    result = wrapper.classify(row)

    assert wrapper.inner.seen == ["a/one"], "stale entry was trusted"
    assert result.task == "sentiment"


# ==========================================================================
# 5. UI filter state: maintained/graveyard filter and its hidden count
# ==========================================================================


def test_5_maintained_filter_hides_unmaintained_and_count_matches(index):
    """The filter hides exactly the unmaintained rows, and says how many."""
    state = atlas.default_state()
    assert state["maintained_only"] is True, "the design opens with it on"

    shown = index.filtered(state)
    assert all(r["maintained"] for r in shown)
    assert len(shown) == 3

    hidden = index.hidden_by_maintained(state)
    assert hidden == 2
    # The number the sidebar prints must equal what is actually withheld.
    assert hidden == len(index.filtered(state, ignore_maintained=True)) - len(shown)

    off = dict(state, maintained_only=False)
    assert len(index.filtered(off)) == 5
    assert index.hidden_by_maintained(off) == 0


def test_5b_hidden_count_respects_the_other_filters(index):
    """Hidden means "hidden from this view", not a global constant."""
    state = dict(atlas.default_state(), tasks=["forecasting"])

    shown = index.filtered(state)
    hidden = index.hidden_by_maintained(state)

    # Two forecasting models: chronos (maintained), forex-lstm (not).
    assert {r["id"] for r in shown} == {"amazon/chronos-t5-small"}
    assert hidden == 1, "hidden count ignored the task filter"


def test_5c_rendered_sidebar_prints_the_true_hidden_count(index):
    from src.ui import shell

    state = atlas.default_state()
    hidden = index.hidden_by_maintained(state)
    html = shell.filter_rail(index, state, hidden)

    assert "2 UNMAINTAINED HIDDEN" in html
    assert str(index.unmaintained_count) == "2"

    off = dict(state, maintained_only=False)
    html_off = shell.filter_rail(index, off, index.hidden_by_maintained(off))
    assert "UNMAINTAINED SHOWN" in html_off


def test_5d_graveyard_action_turns_the_filter_off(index):
    """The design's "SHOW THEM →" button reveals the graveyard."""
    import app as atlas_app

    state = atlas.default_state()
    after = atlas_app.apply_action(state, "graveyard:|nonce")
    assert after["maintained_only"] is False
    assert len(index.filtered(after)) == 5


# ==========================================================================
# 6. Lineage walk terminates on circular base_model references
# ==========================================================================


def test_6_circular_lineage_terminates(index):
    """Two models each declaring the other must not hang the walk."""
    circular = [
        make_row("a/one", base_model="b/two"),
        make_row("b/two", base_model="a/one"),
    ]
    parents = lineage.build_parents(circular)
    children = lineage.build_children(parents)

    # If the guard is missing, these spin forever rather than failing.
    assert lineage.ancestors("a/one", parents) == ["b/two"]
    assert lineage.descendants("a/one", children) == ["b/two"]
    assert lineage.root_of("a/one", parents) == "b/two"


def test_6b_self_referential_base_model_terminates():
    rows = [make_row("a/one", base_model="a/one")]
    parents = lineage.build_parents(rows)
    children = lineage.build_children(parents)

    assert lineage.ancestors("a/one", parents) == []
    assert lineage.descendants("a/one", children) == []


def test_6c_longer_cycle_terminates():
    """A three-model ring: A -> B -> C -> A."""
    rows = [
        make_row("a/one", base_model="b/two"),
        make_row("b/two", base_model="c/three"),
        make_row("c/three", base_model="a/one"),
    ]
    parents = lineage.build_parents(rows)
    children = lineage.build_children(parents)

    chain = lineage.ancestors("a/one", parents)
    assert chain == ["b/two", "c/three"], chain
    assert len(chain) == len(set(chain)), "a node was visited twice"

    kids = lineage.descendants("a/one", children)
    assert sorted(kids) == ["b/two", "c/three"]


def test_6d_cycle_does_not_hang_the_spotlight():
    """`largest_family` walks every node; a cycle must not trap it."""
    rows = [
        make_row("a/one", base_model="b/two"),
        make_row("b/two", base_model="a/one"),
        make_row("root/parent"),
        make_row("kid/one", base_model="root/parent"),
        make_row("kid/two", base_model="root/parent"),
    ]
    root, kids = lineage.largest_family(rows)
    assert root == "root/parent"
    assert sorted(kids) == ["kid/one", "kid/two"]


def test_6e_deep_chain_is_depth_capped():
    """A pathological chain stops at MAX_DEPTH rather than walking forever."""
    rows = [make_row(f"org/m{i}", base_model=f"org/m{i + 1}") for i in range(100)]
    parents = lineage.build_parents(rows)

    chain = lineage.ancestors("org/m0", parents)
    assert len(chain) <= lineage.MAX_DEPTH


def test_6f_drawer_renders_for_a_model_in_a_cycle():
    """The end-to-end path: a cycle must not break rendering either."""
    import pandas as pd
    from src.ui import shell

    circular = [
        make_row("a/one", base_model="b/two"),
        make_row("b/two", base_model="a/one"),
    ]
    built = atlas.Index(pd.DataFrame(circular))
    built.dataset_repo = "x/y"

    html = shell.drawer(built.by_id["a/one"], built)
    assert "a/one" in html
    assert "b/two" in html


# ==========================================================================
# Grounding: a summary the card cannot support is not a summary
# ==========================================================================
#
# "Undocumented" is one of the six headline numbers on the page, and its whole
# job is to count models that never say what they were trained on. A plausible
# guess in that column does not make the number slightly worse -- it inverts
# what it means. Measured on a real run, ~24% of models with a card under 250
# characters still came back with a confident summary.


def test_summary_is_dropped_when_the_card_is_too_thin(monkeypatch):
    """The real failure: a 27-character card described as 'financial text'."""
    answer = json.dumps({
        "relevant": "yes", "task": "sentiment", "asset_class": "general",
        "training_data_summary": "finetuned on financial text",
        "maintained": True, "has_eval": False, "red_flags": [],
    })
    clf, _ = make_classifier([answer], monkeypatch)
    row = make_row("neonbit01/finbert-finetuned-v2", readme="# finbert v2\n")

    result = clf.classify(row)

    assert result.training_data_summary == ""
    assert classify.UNGROUNDED_FLAG in result.red_flags


def test_summary_is_dropped_when_the_card_never_mentions_training_data():
    """A long card about something else cannot ground a training-data claim."""
    row = make_row("a/b", readme="# Model\n\n" + ("Usage instructions. " * 40))
    result = schema.Classification(
        relevant="yes", task="sentiment", asset_class="general",
        training_data_summary="trained on financial news")

    grounded = classify.ground_summary(row, result)
    assert grounded.training_data_summary == ""


def test_grounded_summary_survives():
    """A card that genuinely documents its corpus keeps its summary."""
    card = ("# FinBERT\n\nFinBERT is a pre-trained NLP model to analyse "
            "sentiment of financial text. It is built by further training BERT "
            "in the finance domain, using a large financial corpus. The "
            "Financial PhraseBank dataset by Malo et al. (2014) is used for "
            "fine-tuning, with an additional held-out split for evaluation.")
    row = make_row("ProsusAI/finbert", readme=card)
    result = schema.Classification(
        relevant="yes", task="sentiment", asset_class="general",
        training_data_summary="Financial PhraseBank and a large financial corpus.")

    grounded = classify.ground_summary(row, result)
    assert grounded.training_data_summary
    assert classify.UNGROUNDED_FLAG not in grounded.red_flags


def test_an_empty_summary_is_left_alone():
    row = make_row("a/b", readme="")
    result = schema.Classification(relevant="unclear", training_data_summary="")
    assert classify.ground_summary(row, result).red_flags == []


def test_grounding_is_applied_to_cache_hits_too(tmp_path):
    """The guard is deterministic, so poisoned cache entries self-correct.

    Without this, every summary invented before the guard existed would keep
    being served until someone paid to reclassify the whole index.
    """
    cache_path = tmp_path / "cache.json"
    row = make_row("a/b", readme="# tiny\n")

    poisoned = {
        "relevant": "yes", "task": "sentiment", "asset_class": "general",
        "training_data_summary": "trained on financial text",
        "maintained": True, "has_eval": False, "red_flags": [],
    }
    cache_path.write_text(json.dumps({classify.cache_key(row): poisoned}))

    wrapper = classify.CachedClassifier(CountingClassifier(), cache_path)
    result = wrapper.classify(row)

    assert wrapper.hits == 1, "should still be a cache hit, not a reclassify"
    assert wrapper.inner.seen == []
    assert result.training_data_summary == "", "served an ungrounded summary"