airdef/tools/test_budget_report.py

"""Offline failure-path tests for AirDef's opt-in WASM budget report."""
import hashlib
import json
import os
import tempfile
import unittest
from pathlib import Path
from subprocess import CompletedProcess
from unittest import mock

import airdef


class BudgetReportFailureTests(unittest.TestCase):
    def setUp(self):
        self.temp = tempfile.TemporaryDirectory(prefix="airdef-budget-report-test-")
        self.addCleanup(self.temp.cleanup)
        self.root = Path(self.temp.name)
        self.fork = self.root / "horizon"
        self.mod_root = self.root / "airdef-mod"
        self.artifact = self.mod_root / "target" / "horizon-mod" / "airdef.hmod"
        self.report = self.root / "report.json"

        (self.fork / "crates/horizon-mod-sdk/wit").mkdir(parents=True)
        (self.fork / "crates/horizon-mods/src").mkdir(parents=True)
        (self.mod_root / "tests/budgets").mkdir(parents=True)
        (self.mod_root / "src").mkdir(parents=True)
        self.artifact.parent.mkdir(parents=True)
        (self.fork / "crates/horizon-mod-sdk/wit/horizon.wit").write_text("package test:mods;\n")
        (self.fork / "crates/horizon-mods/src/runtime.rs").write_text(
            "const MEMORY: usize = 64 << 20;\n"
            "const FUEL_INIT: u64 = 2000000000;\n"
            "const FUEL_TICK: u64 = 300000000;\n"
            "const FUEL_FRAME: u64 = 10000000;\n"
            "const FUEL_EVENT: u64 = 1000000000;\n"
            "wasm_component_model(true); consume_fuel(true);\n"
        )
        (self.fork / "Cargo.lock").write_text(
            '[[package]]\nname = "wasmtime"\nversion = "48.0.7"\n'
        )
        (self.mod_root / "src/layout.rs").write_text("pub const MAX_UNITS: usize = 2000;\n")
        (self.mod_root / "tests/budgets/wasm_flat_budget.toml").write_text(
            '[[step]]\ncommand = "airdef test"\n'
        )
        self.artifact.write_bytes(b"minimal synthetic hmod")
        self.diff_sha = hashlib.sha256(b"synthetic diff").hexdigest()
        self.wit_sha = hashlib.sha256((self.fork / "crates/horizon-mod-sdk/wit/horizon.wit").read_bytes()).hexdigest()
        self.identity = {
            "airdef": "airdef-sha",
            "airdef_dirty": False,
            "airdef_diff_sha256": self.diff_sha,
            "engine": "engine-id",
            "engine_pin": "engine-pin",
            "wit_sha256": self.wit_sha,
            "artifact_sha256": hashlib.sha256(self.artifact.read_bytes()).hexdigest(),
        }
        (self.artifact.parent / "build-identity.json").write_text(json.dumps(self.identity))

        self.patches = [
            mock.patch.object(airdef, "ROOT", self.root),
            mock.patch.object(airdef, "FORK", self.fork),
            mock.patch.object(airdef, "engine_commit", return_value="engine-pin"),
            mock.patch.object(airdef, "engine_identity", return_value="engine-id"),
            mock.patch.object(airdef, "fork_api", return_value=(27, [])),
            mock.patch.object(airdef, "hmod", return_value=self.artifact),
            mock.patch.object(airdef, "cargo", return_value="cargo"),
            mock.patch.dict(os.environ, {"AIRDEF_BUDGET_REPORT": str(self.report)}),
            mock.patch.object(airdef.subprocess, "check_output", side_effect=self.fake_git),
        ]
        for patcher in self.patches:
            patcher.start()
            self.addCleanup(patcher.stop)

    @staticmethod
    def fake_git(command, **kwargs):
        if "rev-parse" in command:
            result = "airdef-sha\n"
        elif "status" in command:
            result = ""
        elif "diff" in command:
            result = b"synthetic diff" if kwargs.get("text") is False else "synthetic diff"
        else:
            raise AssertionError(f"unexpected git command: {command}")
        return result

    def read_report(self):
        return json.loads(self.report.read_text())

    def invoke_and_require_failure(self):
        with self.assertRaises(SystemExit):
            airdef.test_budgets()
        self.assertTrue(self.report.is_file(), "failure path must persist its report")
        report = self.read_report()
        self.assertEqual(report["status"], "failed")
        return report

    def test_counts_are_requests_and_derived_from_fixture(self):
        (self.mod_root / "tests/budgets/wasm_flat_budget.toml").write_text(
            '[[step]]\ncommand = "airdef deploy s400"\n'
            '[[step]]\ncommand = "airdef wave kh101 80 4000"\n'
            '[[step]]\nticks = 450\nframes = 225\n'
            '[[step]]\ncommand = "airdef wave mix 40 1500"\n'
        )
        with mock.patch.object(airdef, "mod", side_effect=SystemExit("stop after provenance")):
            counts = self.invoke_and_require_failure()["scenario_counts"]
        self.assertEqual((counts["steps"], counts["air_defense_units"]), (4, 1))
        self.assertEqual((counts["raid_requests"], counts["raid_units_requested"]), (2, 120))
        self.assertEqual((counts["ticks"], counts["frames"]), (450, 225))
        self.assertIsNone(counts["accepted_raids_asserted"])
        self.assertIsNone(counts["peak_airborne_observed"])

    def test_success_reports_only_assertions_of_exact_fixture(self):
        source = Path(__file__).resolve().parent.parent / "airdef-mod" / airdef.BUDGET_SCENARIO
        (self.mod_root / airdef.BUDGET_SCENARIO).write_bytes(source.read_bytes())
        def backend(command, cwd, **kwargs):
            self.assertEqual((Path(cwd) / "tests/wasm_budget.toml").read_bytes(), source.read_bytes())
            self.assertIn("--fuel-profile", command)
            return CompletedProcess(args=command, returncode=0, stdout=(
                "test passed\nfuel airdef tick: calls=2102 max=123456 budget=300000000 failures=0\n"))
        with mock.patch.object(airdef, "mod"), mock.patch.object(airdef.subprocess, "run", side_effect=backend):
            airdef.test_budgets()
        report = self.read_report()
        self.assertEqual(report["scenario"], airdef.BUDGET_SCENARIO)
        self.assertEqual(report["scenario_render_world"], "flat at y=64")
        counts = report["scenario_counts"]
        self.assertEqual((counts["accepted_raids_asserted"], counts["snapshot_airborne_asserted"], counts["snapshot_tick"]), (4, 184, 2102))
        self.assertIsNone(counts["peak_airborne_observed"])
        self.assertEqual(report["measured_fuel"]["callbacks"], [{
            "mod": "airdef", "callback": "tick", "calls": 2102,
            "max_used": 123456, "budget": 300000000, "failures": 0}])

    def test_success_without_requested_measurements_is_not_a_budget_proof(self):
        backend = mock.Mock(return_value=CompletedProcess(args=[], returncode=0, stdout="test passed\n"))
        with mock.patch.object(airdef, "mod"), mock.patch.object(airdef.subprocess, "run", backend):
            report = self.invoke_and_require_failure()
        self.assertFalse(report["measured_fuel"]["available"])
        self.assertIn("did not return", report["error"])

    def test_early_system_exit_writes_report_and_skips_backend(self):
        backend = mock.Mock()
        with mock.patch.object(airdef, "mod", side_effect=SystemExit("synthetic build failure")), \
             mock.patch.object(airdef.subprocess, "run", backend):
            report = self.invoke_and_require_failure()

        self.assertIn("SystemExit", report["error"])
        backend.assert_not_called()

    def test_stale_artifact_identity_blocks_backend(self):
        stale = dict(self.identity, engine_pin="old-engine-pin")
        (self.artifact.parent / "build-identity.json").write_text(json.dumps(stale))
        backend = mock.Mock()
        with mock.patch.object(airdef, "mod"), \
             mock.patch.object(airdef.subprocess, "run", backend):
            report = self.invoke_and_require_failure()

        self.assertFalse(report["build_identity"]["matches"])
        self.assertFalse(report["build_identity"]["checks"]["engine_pin"])
        self.assertIn("stale", report["error"])
        backend.assert_not_called()

    def test_runtime_fuel_trap_persists_nonzero_exit_and_output(self):
        backend = mock.Mock(return_value=CompletedProcess(
            args=["cargo", "run"], returncode=137, stdout=("error: all fuel consumed (fuel trap)\n"
                "fuel airdef tick: calls=4 max=300000000 budget=300000000 failures=1\n")
        ))
        with mock.patch.object(airdef, "mod"), \
             mock.patch.object(airdef.subprocess, "run", backend):
            report = self.invoke_and_require_failure()

        backend.assert_called_once()
        self.assertEqual(report["exit_code"], 137)
        self.assertIn("fuel trap", report["output"])
        self.assertEqual(report["measured_fuel"]["callbacks"][0]["failures"], 1)


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