tests/test_rules_engine.py

"""Tests for ``RulesEngine``.

These tests cover the three things the engine is responsible for:

1. **Hard prerequisites** — actions the agent cannot perform until earlier
   pipeline milestones are unlocked (e.g. submit-claim before
   estimate-significance is rejected).
2. **Soft violations** — invalid params, redundancy, out-of-window.
3. **Resource gating** — once the budget / time / luminosity is exhausted
   the engine refuses further actions.
"""

from __future__ import annotations

import pytest

from models import (
    ActionType,
    DiscoveryClaim,
    ExperimentAction,
)
from server.rules.engine import RulesEngine, ViolationCode
from server.tasks.scenarios import sample_scenario


@pytest.fixture
def fresh_state():
    """A fresh latent state for the easy diphoton scenario."""
    sc = sample_scenario(name="easy_diphoton_160", seed=7)
    return sc.fresh_latent()


@pytest.fixture
def rules():
    return RulesEngine(mass_search_window_gev=(80.0, 300.0))


# ── Prerequisites ────────────────────────────────────────────────────────


def test_collect_collisions_blocked_without_setup(rules, fresh_state):
    action = ExperimentAction(action_type=ActionType.COLLECT_COLLISIONS)
    result = rules.validate(action, fresh_state)
    assert not result.allowed
    assert ViolationCode.PREREQ_MISSING in result.violations


def test_fit_resonance_blocked_without_histogram(rules, fresh_state):
    action = ExperimentAction(action_type=ActionType.FIT_RESONANCE)
    result = rules.validate(action, fresh_state)
    assert not result.allowed
    assert ViolationCode.PREREQ_MISSING in result.violations


def test_submit_claim_blocked_without_significance(rules, fresh_state):
    fresh_state.progress.resonance_fitted = True  # pretend we got that far
    action = ExperimentAction(
        action_type=ActionType.SUBMIT_DISCOVERY_CLAIM,
        parameters={"claim": {"mass_estimate_gev": 125.0, "significance_sigma": 5.0}},
    )
    result = rules.validate(action, fresh_state)
    assert not result.allowed
    # the missing-significance prereq is the dominant failure
    assert ViolationCode.PREREQ_MISSING in result.violations


# ── Resource gating ──────────────────────────────────────────────────────


def test_budget_exhausted_blocks_everything(rules, fresh_state):
    fresh_state.resources.budget_used_musd = fresh_state.resources.budget_total_musd
    action = ExperimentAction(action_type=ActionType.CONFIGURE_BEAM)
    result = rules.validate(action, fresh_state)
    assert not result.allowed
    assert ViolationCode.BUDGET_EXHAUSTED in result.violations


def test_time_exhausted_blocks_everything(rules, fresh_state):
    fresh_state.resources.time_used_days = fresh_state.resources.time_limit_days
    action = ExperimentAction(action_type=ActionType.CONFIGURE_BEAM)
    result = rules.validate(action, fresh_state)
    assert not result.allowed
    assert ViolationCode.TIME_EXHAUSTED in result.violations


def test_luminosity_exhaustion_only_blocks_daq(rules, fresh_state):
    fresh_state.resources.luminosity_used_fb = fresh_state.resources.luminosity_total_fb
    blocked = ExperimentAction(action_type=ActionType.COLLECT_COLLISIONS)
    allowed = ExperimentAction(action_type=ActionType.CONFIGURE_BEAM)
    assert not rules.validate(blocked, fresh_state).allowed
    assert rules.validate(allowed, fresh_state).allowed


# ── Soft violations ──────────────────────────────────────────────────────


def test_unknown_channel_is_soft_violation(rules, fresh_state):
    action = ExperimentAction(
        action_type=ActionType.SELECT_CHANNEL,
        parameters={"channel": "purple_quark"},
    )
    result = rules.validate(action, fresh_state)
    assert result.allowed  # soft
    assert ViolationCode.INVALID_PARAMS in result.soft_violations


def test_redundant_beam_config_is_soft_violation(rules, fresh_state):
    fresh_state.progress.beam_configured = True
    action = ExperimentAction(action_type=ActionType.CONFIGURE_BEAM)
    result = rules.validate(action, fresh_state)
    assert result.allowed
    assert ViolationCode.REDUNDANT in result.soft_violations


def test_inverted_mass_window_is_soft_violation(rules, fresh_state):
    action = ExperimentAction(
        action_type=ActionType.BUILD_INVARIANT_MASS,
        parameters={"mass_window_gev": [200.0, 100.0]},
    )
    result = rules.validate(action, fresh_state)
    # rules engine flags hi<=lo as soft INVALID_PARAMS
    assert ViolationCode.INVALID_PARAMS in result.soft_violations


def test_out_of_window_histogram_is_soft_violation(rules, fresh_state):
    action = ExperimentAction(
        action_type=ActionType.BUILD_INVARIANT_MASS,
        parameters={"mass_window_gev": [10000.0, 20000.0]},
    )
    result = rules.validate(action, fresh_state)
    assert ViolationCode.OUT_OF_WINDOW in result.soft_violations


def test_claim_missing_mass_is_invalid(rules, fresh_state):
    fresh_state.progress.resonance_fitted = True
    fresh_state.progress.significance_estimated = True
    action = ExperimentAction(
        action_type=ActionType.SUBMIT_DISCOVERY_CLAIM,
        parameters={"claim": {"significance_sigma": 5.0}},
    )
    result = rules.validate(action, fresh_state)
    assert not result.allowed
    assert ViolationCode.INVALID_CLAIM in result.violations


def test_well_formed_claim_passes_rules(rules, fresh_state):
    fresh_state.progress.resonance_fitted = True
    fresh_state.progress.significance_estimated = True
    action = ExperimentAction(
        action_type=ActionType.SUBMIT_DISCOVERY_CLAIM,
        parameters={
            "claim": {
                "mass_estimate_gev": 160.0,  # inside [80, 300]
                "significance_sigma": 5.2,
            }
        },
    )
    result = rules.validate(action, fresh_state)
    assert result.allowed
    assert not result.violations