SystemSimulationApp/tests/test_max_step_matrix.py

from __future__ import annotations

from contextlib import redirect_stdout
from io import StringIO
from pathlib import Path
import unittest
from unittest.mock import patch

from app.simulation.benchmark_regression import (
    DEFAULT_MANIFEST_PATH,
    RegressionCaseRequest,
    RegressionManifestError,
)
from app.simulation.max_step_matrix import main, run_max_step_matrix


def _fake_completed(
    request: RegressionCaseRequest,
    *,
    prefix_bias: float = 0.0,
) -> dict[str, object]:
    signal_times = [value for value in (0.04, 0.8) if value <= request.stop_time]
    checkpoints = [
        {
            "requestedTime": checkpoint,
            "actualTime": checkpoint,
            "available": True,
            "stateValues": {
                "state.a": checkpoint + prefix_bias,
                "state.b": 2.0 * checkpoint,
            },
        }
        for checkpoint in request.checkpoint_times
    ]
    event_trace = {
        "signalEventTimes": signal_times,
        "stateTransitionCount": 0,
        "mechanicalTransitionTimes": [],
        "mechanicalTransitionTimesAvailable": True,
    }
    return {
        "outcome": "completed",
        "orchestrationWallSeconds": 0.01,
        "worker": {
            "outcome": "completed",
            "wallSeconds": 0.01,
            "summary": {
                "success": True,
                "status": "completed",
                "simulatedUntil": request.stop_time,
                "checkpoints": checkpoints,
                "physicalContract": {
                    "schemaVersion": 1,
                    "projectionCategories": ["state"],
                    "checkpoints": checkpoints,
                    "eventTrace": event_trace,
                },
                "eventTrace": event_trace,
                "diagnostics": {
                    "pressureFlow": {"maxScaledResidual": 1.0e-12},
                    "integration": {
                        "totals": {
                            "nfev": int(100 * request.stop_time / request.max_step),
                            "njev": 4,
                            "nlu": 8,
                            "acceptedStepCount": int(
                                request.stop_time / request.max_step
                            ),
                            "solverStartCount": len(signal_times) + 1,
                            "stateTransitionCount": 0,
                            "recoverableRetryCount": 0,
                        }
                    },
                },
            },
        },
    }


def _fake_failure(request: RegressionCaseRequest) -> dict[str, object]:
    return {
        "outcome": "failed",
        "worker": {
            "outcome": "failed",
            "summary": {
                "success": False,
                "status": "failed",
                "simulatedUntil": 0.75 * request.stop_time,
            },
        },
    }


class MaxStepMatrixTests(unittest.TestCase):
    def test_single_cell_passes_when_no_comparison_is_required(self) -> None:
        report = run_max_step_matrix(
            DEFAULT_MANIFEST_PATH,
            horizon_case_ids=("1s",),
            max_steps=(0.02,),
            expected_projection_count=2,
            case_executor=_fake_completed,
        )

        comparisons = report["comparisons"]
        self.assertEqual(comparisons["sameHorizonAcrossMaxSteps"], [])
        self.assertEqual(comparisons["sameMaxStepAcrossHorizons"], [])
        self.assertEqual(comparisons["evaluatedCount"], 0)
        self.assertTrue(comparisons["passed"])
        self.assertTrue(report["acceptance"]["passed"])

        with patch(
            "app.simulation.max_step_matrix.run_max_step_matrix",
            return_value=report,
        ), redirect_stdout(StringIO()):
            exit_code = main(["--horizon", "1s", "--max-step", "0.02"])
        self.assertEqual(exit_code, 0)

    def test_multiple_cells_still_fail_when_comparison_is_unavailable(self) -> None:
        def execute(request: RegressionCaseRequest) -> dict[str, object]:
            if request.max_step == 0.05:
                return _fake_failure(request)
            return _fake_completed(request)

        report = run_max_step_matrix(
            DEFAULT_MANIFEST_PATH,
            horizon_case_ids=("1s",),
            max_steps=(0.01, 0.05),
            expected_projection_count=2,
            case_executor=execute,
        )

        comparisons = report["comparisons"]
        records = comparisons["sameHorizonAcrossMaxSteps"]
        self.assertEqual(len(records), 1)
        self.assertFalse(records[0]["evaluated"])
        self.assertEqual(records[0]["reason"], "oneOrBothCasesDidNotComplete")
        self.assertFalse(comparisons["passed"])
        self.assertFalse(report["acceptance"]["passed"])

        with patch(
            "app.simulation.max_step_matrix.run_max_step_matrix",
            return_value=report,
        ), redirect_stdout(StringIO()):
            exit_code = main(
                [
                    "--horizon",
                    "1s",
                    "--max-step",
                    "0.01",
                    "--max-step",
                    "0.05",
                ]
            )
        self.assertEqual(exit_code, 1)

    def test_custom_horizon_and_common_prefix_matrix(self) -> None:
        source = Path("tests/data/test-mql-8.xml")
        source_before = source.read_bytes()
        requests: list[RegressionCaseRequest] = []

        def execute(request: RegressionCaseRequest) -> dict[str, object]:
            requests.append(request)
            return _fake_completed(request)

        report = run_max_step_matrix(
            DEFAULT_MANIFEST_PATH,
            horizon_case_ids=("1s", 2.0),
            max_steps=(0.01, 0.05),
            additional_checkpoint_times=(1.08, 1.10),
            expected_projection_count=2,
            case_executor=execute,
        )

        self.assertEqual(source.read_bytes(), source_before)
        self.assertEqual(len(requests), 4)
        self.assertEqual(
            [(request.stop_time, request.max_step) for request in requests],
            [(1.0, 0.01), (1.0, 0.05), (2.0, 0.01), (2.0, 0.05)],
        )
        custom_requests = [request for request in requests if request.stop_time == 2.0]
        self.assertEqual(
            custom_requests[0].checkpoint_times,
            (0.0, 0.04, 0.8, 1.0, 1.08, 1.1, 2.0),
        )
        self.assertEqual(
            report["configuration"]["horizons"][1]["horizonCaseId"],
            "2s-custom",
        )
        self.assertEqual(
            len(report["comparisons"]["sameHorizonAcrossMaxSteps"]), 2
        )
        self.assertEqual(
            len(report["comparisons"]["sameMaxStepAcrossHorizons"]), 2
        )
        self.assertTrue(report["comparisons"]["passed"])
        self.assertTrue(report["acceptance"]["passed"])

    def test_failed_tier_defers_every_later_cell(self) -> None:
        requests: list[RegressionCaseRequest] = []

        def execute(request: RegressionCaseRequest) -> dict[str, object]:
            requests.append(request)
            if request.max_step == 0.05:
                return _fake_failure(request)
            return _fake_completed(request)

        report = run_max_step_matrix(
            DEFAULT_MANIFEST_PATH,
            horizon_case_ids=("1s", "5s", "10s"),
            max_steps=(0.01, 0.05),
            expected_projection_count=2,
            case_executor=execute,
        )

        self.assertEqual(len(requests), 2)
        self.assertEqual(report["acceptance"]["deferredCellCount"], 4)
        self.assertEqual(
            [case["outcome"] for case in report["cases"][2:]],
            ["deferred"] * 4,
        )
        self.assertEqual(
            report["tierDecisions"][1]["reason"], "previousTierDidNotPass"
        )

    def test_tstop_dependent_prefix_is_reported(self) -> None:
        def execute(request: RegressionCaseRequest) -> dict[str, object]:
            return _fake_completed(
                request,
                prefix_bias=0.1 if request.stop_time > 1.0 else 0.0,
            )

        report = run_max_step_matrix(
            DEFAULT_MANIFEST_PATH,
            horizon_case_ids=("1s", 2.0),
            max_steps=(0.01, 0.05),
            expected_projection_count=2,
            stop_after_failed_tier=False,
            case_executor=execute,
        )

        prefix_comparisons = report["comparisons"][
            "sameMaxStepAcrossHorizons"
        ]
        self.assertEqual(len(prefix_comparisons), 2)
        self.assertTrue(all(item["evaluated"] for item in prefix_comparisons))
        self.assertTrue(all(not item["passed"] for item in prefix_comparisons))
        self.assertGreater(
            prefix_comparisons[0]["stateProjection"]["valueMismatchCount"], 0
        )

    def test_cross_horizon_comparison_excludes_short_terminal(self) -> None:
        def execute(request: RegressionCaseRequest) -> dict[str, object]:
            result = _fake_completed(request)
            if request.stop_time <= 1.0:
                checkpoints = result["worker"]["summary"]["physicalContract"][
                    "checkpoints"
                ]
                for checkpoint in checkpoints:
                    if checkpoint["requestedTime"] == request.stop_time:
                        checkpoint["stateValues"]["state.a"] += 1.0
            return result

        report = run_max_step_matrix(
            DEFAULT_MANIFEST_PATH,
            horizon_case_ids=("1s", 2.0),
            max_steps=(0.01,),
            expected_projection_count=2,
            case_executor=execute,
        )

        comparison = report["comparisons"]["sameMaxStepAcrossHorizons"][0]
        self.assertTrue(comparison["strictPrefix"])
        self.assertTrue(comparison["passed"])
        self.assertNotIn(
            1.0,
            comparison["stateProjection"]["commonCheckpointTimes"],
        )

    def test_same_horizon_comparison_keeps_terminal_checkpoint(self) -> None:
        report = run_max_step_matrix(
            DEFAULT_MANIFEST_PATH,
            horizon_case_ids=("1s",),
            max_steps=(0.01, 0.05),
            expected_projection_count=2,
            case_executor=_fake_completed,
        )

        comparison = report["comparisons"]["sameHorizonAcrossMaxSteps"][0]
        self.assertFalse(comparison["strictPrefix"])
        self.assertIn(
            1.0,
            comparison["stateProjection"]["commonCheckpointTimes"],
        )

    def test_horizons_must_increase_and_timeout_override_is_bounded(self) -> None:
        with self.assertRaisesRegex(
            RegressionManifestError, "strictly increasing"
        ):
            run_max_step_matrix(
                DEFAULT_MANIFEST_PATH,
                horizon_case_ids=(2.0, "1s"),
                max_steps=(0.01, 0.05),
                expected_projection_count=2,
                case_executor=_fake_completed,
            )

    def test_off_grid_checkpoint_is_rejected_before_execution(self) -> None:
        with self.assertRaisesRegex(
            RegressionManifestError,
            "not represented by the .* output grid",
        ):
            run_max_step_matrix(
                DEFAULT_MANIFEST_PATH,
                horizon_case_ids=("1s",),
                max_steps=(0.01,),
                additional_checkpoint_times=(0.0489,),
                expected_projection_count=2,
                case_executor=_fake_completed,
            )
        with self.assertRaisesRegex(RegressionManifestError, "must exceed"):
            run_max_step_matrix(
                DEFAULT_MANIFEST_PATH,
                horizon_case_ids=(2.0,),
                max_steps=(0.01, 0.05),
                soft_timeout_seconds=20.0,
                hard_timeout_seconds=10.0,
                expected_projection_count=2,
                case_executor=_fake_completed,
            )


if __name__ == "__main__":
    unittest.main()