yau-plant-assistant/api/tests/test_contracts.py

"""Contract tests — the safety rules, exercised without an API key.

These run in CI and locally with `pytest api/tests`. They need no network, no
database and no model, because the contracts are pure Python. That is the point
of putting the safety rules there.
"""

from datetime import date

import pytest

from contracts import (
    AdvisoryAnswer,
    Citation,
    ContractViolation,
    Evidence,
    HistoricalAnswer,
    ProceduralAnswer,
    ProcedureIdentity,
    QuestionClass,
    TimeWindow,
)

WINDOW = TimeWindow(
    start="2026-08-13T09:00:00",
    end="2026-08-20T09:00:00",
    timezone="Australia/Sydney",
    description="rolling 7 days to 2026-08-20 09:00 AEST",
)

PROCEDURE = ProcedureIdentity(
    doc_number="WRPS-OPS-014",
    title="Pump Interlock Lifting",
    revision="3",
    effective_date=date(2025, 11, 3),
    authorising_role="Operations Supervisor",
    controlled_copy_location="Document control, WRPS station office",
)

CITATION = Citation(
    doc_number="WRPS-OPS-014",
    title="Pump Interlock Lifting",
    revision="3",
    effective_date=date(2025, 11, 3),
    source_file="procedures/WRPS-OPS-014.pdf",
)


def procedural(answer: str, **kwargs) -> ProceduralAnswer:
    return ProceduralAnswer(
        question="How do I lift the interlock on Pump 02?",
        answer=answer,
        procedure=kwargs.pop("procedure", PROCEDURE),
        citations=kwargs.pop("citations", [CITATION]),
        **kwargs,
    )


# --- Procedural: never synthesise steps -------------------------------------


def test_procedural_identification_is_allowed():
    result = procedural(
        "The governing procedure is WRPS-OPS-014 rev 3, effective 3 November "
        "2025, authorised by the Operations Supervisor. Work from the "
        "controlled copy held by document control.",
        prerequisites_verbatim=[
            "The unit shall be confirmed stopped and isolated before any "
            "interlock is lifted."
        ],
    )
    assert result.steps_provided is False


@pytest.mark.parametrize(
    "text",
    [
        "To lift the interlock, first navigate to the pump faceplate and set the mode to manual.",
        "Follow these steps: open the maintenance screen and press override.",
        "You should set the station mode to off, then bypass the interlock.",
        "Here is how to reset the trip on PU-302.",
    ],
)
def test_procedural_rejects_instructions(text):
    with pytest.raises(ContractViolation) as caught:
        procedural(text)
    assert caught.value.rule in {"synthesised_steps", "step_sequence_emitted"}


def test_procedural_rejects_numbered_step_sequence():
    with pytest.raises(ContractViolation) as caught:
        procedural("1. Stop the pump.\n2. Isolate the supply.\n3. Lift the interlock.")
    assert caught.value.rule in {"synthesised_steps", "step_sequence_emitted"}


def test_procedural_with_no_procedure_must_say_so():
    with pytest.raises(ContractViolation) as caught:
        procedural("The interlock is on the pump control block.", procedure=None)
    assert caught.value.rule == "no_procedure_no_refusal"


# --- Advisory: evidence and deferral, never a recommended value -------------


def advisory(answer: str, **kwargs) -> AdvisoryAnswer:
    return AdvisoryAnswer(
        question="What is the best discharge rate to draw the well down?",
        answer=answer,
        evidence=kwargs.pop(
            "evidence",
            [
                Evidence(
                    description="Mean discharge rate across two-pump operations",
                    metric="operations.avg_discharge_rate",
                    value=864.0,
                    unit="m3/h",
                    sample_size=118,
                    time_window=WINDOW,
                )
            ],
        ),
        deferral=kwargs.pop(
            "deferral",
            "The operating rate is a decision for a competent person with sight "
            "of current equipment condition and concurrent operations.",
        ),
        **kwargs,
    )


def test_advisory_evidence_and_deferral_is_allowed():
    result = advisory(
        "Over the last 30 days the station ran 118 two-pump operations at a mean "
        "discharge of 864 m3/h. Twelve of those reached the high level alarm and "
        "none spilled. The documented station capacity is 1296 m3/h."
    )
    assert result.recommendation_given is False


@pytest.mark.parametrize(
    "text",
    [
        "I recommend running at 900 m3/h.",
        "The optimal flowrate is around 864 m3/h.",
        "You should set the rate to 900 m3/h to stay below the alarm.",
        "Aim for 850 m3/h and the well will not reach the weir.",
    ],
)
def test_advisory_rejects_recommendations(text):
    with pytest.raises(ContractViolation) as caught:
        advisory(text)
    assert caught.value.rule == "recommendation_given"


def test_advisory_requires_a_deferral():
    with pytest.raises(ContractViolation) as caught:
        advisory("The station has run between 432 and 1296 m3/h.", deferral="  ")
    assert caught.value.rule == "missing_deferral"


# --- Historical: zero rows means say so -------------------------------------


def historical(answer: str, rows: list[dict]) -> HistoricalAnswer:
    return HistoricalAnswer(
        question="How many high level alarms last week?",
        answer=answer,
        query={"measures": ["alarm_activity.alarm_count"]},
        row_count=len(rows),
        rows=rows,
        time_window=WINDOW,
    )


def test_historical_with_rows_is_allowed():
    result = historical(
        "The wet well high level alarm activated 6 times in the rolling 7 days "
        "to 2026-08-20 09:00 AEST.",
        [{"alarm_activity.alarm_count": 6}],
    )
    assert result.row_count == 1


def test_historical_zero_rows_must_say_no_records():
    with pytest.raises(ContractViolation) as caught:
        historical("The high level alarm activated 6 times last week.", [])
    assert caught.value.rule in {"zero_rows_not_declared", "figure_without_data"}


def test_historical_zero_rows_saying_so_is_allowed():
    result = historical(
        "No records were found for the wet well high level alarm in the rolling "
        "7 days to 2026-08-20 09:00 AEST.",
        [],
    )
    assert result.row_count == 0


def test_historical_row_count_must_match():
    with pytest.raises(ContractViolation) as caught:
        HistoricalAnswer(
            question="q",
            answer="No records were found.",
            query={},
            row_count=5,
            rows=[],
            time_window=WINDOW,
        )
    assert caught.value.rule == "row_count_mismatch"


# --- Citations ---------------------------------------------------------------


def test_superseded_citation_is_rejected():
    with pytest.raises(ContractViolation) as caught:
        Citation(
            doc_number="WRPS-OPS-014",
            title="Pump Interlock Lifting",
            revision="2",
            effective_date=date(2023, 5, 1),
            source_file="procedures/WRPS-OPS-014-rev2.pdf",
            superseded=True,
        )
    assert caught.value.rule == "superseded_citation"