chemenu/tools/chemenu/tests/test_run_budget.py

import json

import pytest
import typer

from chemenu.commands import run_budget


@pytest.fixture(autouse=True)
def isolated_state(tmp_path, monkeypatch):
    """Point the budget state file at a scratch location and pin the session
    id, so tests never touch the real .wikitool_session/ dir or bleed into
    the actual calling shell's session."""
    state_file = tmp_path / "budget.json"
    monkeypatch.setattr(run_budget, "STATE_DIR", tmp_path)
    monkeypatch.setattr(run_budget, "STATE_FILE", state_file)
    monkeypatch.setattr(run_budget, "LOCK_FILE", tmp_path / "budget.lock")
    monkeypatch.setenv("WIKITOOL_SESSION_ID", "test-session")
    return state_file


def test_default_thresholds_match_the_contract():
    assert run_budget.DEFAULT_CALL_LIMIT == 60
    assert run_budget.DEFAULT_LOOP_WINDOW == 3


def test_calls_below_limit_pass_silently():
    for i in range(5):
        run_budget.record_and_check("new", ["entity", "--name", f"Page{i}"], override=False)
    state = run_budget._load_state()
    assert state["test-session"]["count"] == 5


def test_call_limit_trips_past_threshold():
    limit = run_budget.DEFAULT_CALL_LIMIT
    for i in range(limit):
        run_budget.record_and_check("new", ["entity", "--name", f"Page{i}"], override=False)
    with pytest.raises(typer.Exit):
        run_budget.record_and_check("new", ["entity", "--name", "OneTooMany"], override=False)


def test_call_limit_message_mentions_contract_and_override():
    message = run_budget.call_limit_message(61, 60)
    assert "Iteration and Cost Limits" in message
    assert "--override-budget" in message
    assert "61" in message and "60" in message


def test_record_and_check_reports_whether_it_charged():
    assert run_budget.record_and_check("new", ["entity"], override=False) is True
    assert run_budget.record_and_check("search", ["anything"], override=False) is False


def test_refund_gives_back_the_slot_but_keeps_the_history():
    """A declined call did not iterate on the wiki, so it costs nothing - but
    the loop-breaker still has to see that it happened."""
    for i in range(3):
        run_budget.record_and_check("new", ["entity", "--name", f"Page{i}"], override=False)
    run_budget.record_and_check("new", ["source", "--set", "raw_files=nope"], override=False)
    run_budget.refund()
    entry = run_budget._load_state()["test-session"]
    assert entry["count"] == 3
    assert entry["recent"][-1] == "new source --set raw_files=nope"


def test_refund_never_drives_the_counter_negative():
    run_budget.refund()
    run_budget.record_and_check("new", ["entity"], override=False)
    run_budget.refund()
    run_budget.refund()
    assert run_budget._load_state()["test-session"]["count"] == 0


def test_three_identical_declined_calls_still_trip_the_loop_breaker():
    """The counter is refunded, the history is not - which is what makes the
    loop-breaker the right instrument for a repeated broken invocation."""
    for _ in range(3):
        run_budget.record_and_check("new", ["source", "--set", "raw_files=nope"], override=False)
        run_budget.refund()
    with pytest.raises(typer.Exit):
        run_budget.record_and_check("new", ["source", "--set", "raw_files=nope"], override=False)


def test_refused_call_is_not_counted():
    """A refused call never ran, so it must not inflate the number quoted
    back to the user on the next attempt."""
    limit = run_budget.DEFAULT_CALL_LIMIT
    for i in range(limit):
        run_budget.record_and_check("new", ["entity", "--name", f"Page{i}"], override=False)
    for _ in range(3):
        with pytest.raises(typer.Exit):
            run_budget.record_and_check("new", ["entity", "--name", "Blocked"], override=False)
    assert run_budget._load_state()["test-session"]["count"] == limit


def test_stale_sessions_are_pruned_on_write():
    now = 1_000_000.0
    state = {
        "fresh": {"count": 1, "recent": [], "last_seen": now - 60},
        "stale": {"count": 1, "recent": [], "last_seen": now - run_budget.SESSION_TTL_SECONDS - 1},
        "legacy-no-timestamp": {"count": 1, "recent": []},
    }
    pruned = run_budget.prune_state(state, now)
    assert set(pruned) == {"fresh", "legacy-no-timestamp"}


def test_recorded_call_stamps_last_seen():
    run_budget.record_and_check("lint", [], override=False)
    assert "last_seen" in run_budget._load_state()["test-session"]


def test_loop_breaker_trips_on_identical_repeats():
    args = ["add", "--a", "X", "--b", "Y"]
    for _ in range(3):
        run_budget.record_and_check("xref", args, override=False)
    with pytest.raises(typer.Exit):
        run_budget.record_and_check("xref", args, override=False)


def test_loop_breaker_does_not_trip_on_varying_calls():
    for i in range(5):
        run_budget.record_and_check("xref", ["add", "--a", f"X{i}", "--b", "Y"], override=False)
    # No exception raised - varying args are not a loop.


def test_loop_breaker_message_mentions_signature():
    message = run_budget.loop_breaker_message("xref add --a X --b Y", 3)
    assert "xref add --a X --b Y" in message
    assert "Loop-Breaker" in message


def test_override_bypasses_both_gates():
    args = ["add", "--a", "X", "--b", "Y"]
    for _ in range(40):
        run_budget.record_and_check("xref", args, override=True)
    state = run_budget._load_state()
    assert state["test-session"]["count"] == 40


def test_save_state_writes_atomically_and_leaves_no_tmp_file(isolated_state):
    run_budget._save_state({"test-session": {"count": 1, "recent": []}})
    assert isolated_state.exists()
    tmp_file = isolated_state.with_suffix(isolated_state.suffix + ".tmp")
    assert not tmp_file.exists()
    assert run_budget._load_state()["test-session"]["count"] == 1


def test_state_lock_can_be_acquired_and_released_sequentially():
    """Smoke test for the cross-process lock: acquiring and releasing it must
    not raise, and a later locked call still lands normally (flock locks the
    *open file description*, so nesting two separate acquisitions in one
    process would deadlock - this deliberately tests sequential use only)."""
    with run_budget._state_lock():
        pass
    run_budget.record_and_check("lint", [], override=False)
    assert run_budget._load_state()["test-session"]["count"] == 1


def test_budget_status_is_never_gated():
    """Reading the gate must stay possible after the gate trips - that report is
    what the agent owes the user."""
    for _ in range(50):
        run_budget.record_and_check("budget", ["status"], override=False)
    state = run_budget._load_state()
    assert "test-session" not in state


def test_budget_reset_is_counted_like_any_other_call():
    """`reset` clears the counter, so exempting it would let a session step
    around the gate by resetting first."""
    run_budget.record_and_check("budget", ["reset"], override=False)
    assert run_budget._load_state()["test-session"]["count"] == 1


def test_search_is_never_gated():
    """Retrieval is reading, not iterating. Charging for a search would tax the
    one habit that lowers token cost - looking before reading."""
    for _ in range(50):
        run_budget.record_and_check("search", ["Longhorn"], override=False)
    assert "test-session" not in run_budget._load_state()


def test_search_exemption_survives_a_query_that_looks_like_a_subcommand():
    """`is_exempt` reads args[0] as a subcommand for grouped commands; for
    `search` that slot holds the user's query, so the exemption has to be
    command-level or it depends on what was searched for."""
    assert run_budget.is_exempt("search", ["status"])
    assert run_budget.is_exempt("search", ["anything at all"])
    assert not run_budget.is_exempt("publish", ["--message", "x"])


def test_reset_command_requires_yes():
    run_budget.record_and_check("new", ["entity", "--name", "X"], override=False)
    with pytest.raises(typer.Exit):
        run_budget.reset_command(all_sessions=False, yes=False)
    assert "test-session" in run_budget._load_state()


def test_reset_command_clears_current_session():
    run_budget.record_and_check("new", ["entity", "--name", "X"], override=False)
    run_budget.reset_command(all_sessions=False, yes=True)
    state = run_budget._load_state()
    assert "test-session" not in state


def test_reset_all_clears_state_file(isolated_state):
    run_budget.record_and_check("new", ["entity", "--name", "X"], override=False)
    assert isolated_state.exists()
    run_budget.reset_command(all_sessions=True, yes=True)
    assert not isolated_state.exists()


def test_status_command_reports_count(capsys):
    run_budget.record_and_check("new", ["entity", "--name", "X"], override=False)
    run_budget.status_command()
    out = capsys.readouterr().out
    assert "Calls so far: 1" in out