bundle/v0/tests/test_reconciled_analysis.py

from __future__ import annotations

import copy
import hashlib
import json
from pathlib import Path

import pytest

from loss_aware_dro_repro.batch_control import DENIED_AUTHORITY
from loss_aware_dro_repro.core import ContractError, canonical_bytes, sha256_value
from loss_aware_dro_repro.reconciled_analysis import validate_reconciled_bundle


def _write_json(path: Path, value: dict) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(json.dumps(value, indent=2, sort_keys=True) + "\n", encoding="utf-8")


def _file_hash(path: Path) -> str:
    return hashlib.sha256(path.read_bytes()).hexdigest()


def _fixture(
    tmp_path: Path,
    *,
    rows_override: list[dict] | None = None,
) -> tuple[Path, Path, dict]:
    root = tmp_path / "attempt"
    outputs = root / "outputs"
    manifest_hash = "sha256:" + "1" * 64
    tasks = []
    rows = []
    for index, task_id in enumerate(("suite/d000/r00/n010", "suite/d001/r00/n010")):
        run_identity = "sha256:" + hashlib.sha256(task_id.encode()).hexdigest()
        task_hash = "sha256:" + hashlib.sha256((task_id + "task").encode()).hexdigest()
        output_relpath = f"tasks/task-{index}/{run_identity.removeprefix('sha256:')}"
        tasks.append(
            {
                "task_id": task_id,
                "task_hash": task_hash,
                "scientific_run_identity": run_identity,
                "output_relpath": output_relpath,
                "execution_config_hash": "sha256:" + "2" * 64,
            }
        )
        rows.append(
            {
                "task_id": task_id,
                "task_hash": task_hash,
                "run_identity": run_identity,
                "plan_hash": "sha256:" + "3" * 64,
                "dataset": {
                    "suite": "suite",
                    "distribution_id": index,
                    "replicate": 0,
                    "sample_size": 10,
                },
                "solver": {"status": "optimal"},
                "optimization": {
                    "iterations": 1,
                    "initial_stationarity": 2.0,
                    "final_stationarity": 1.0,
                    "stopping": {
                        "reason": "paper_total_penalized_objective_improvement_below_tolerance"
                    },
                },
                "metrics": {
                    "relative_worst_case_improvement": 0.1,
                    "relative_oos_improvement": 0.1,
                    "coverage_final": True,
                },
            }
        )
    manifest = {
        "manifest_hash": manifest_hash,
        "tasks": tasks,
    }
    manifest_path = tmp_path / "source-manifest.json"
    _write_json(manifest_path, manifest)

    recovered = tasks[0]
    failure_body = {
        "schema_version": 1,
        "event": "failed",
        "manifest_hash": manifest_hash,
        "task_id": recovered["task_id"],
        "scientific_run_identity": recovered["scientific_run_identity"],
        "reason": "future_exception:OSError:[Errno 36] Resource deadlock avoided",
        "returncode": None,
        "stdout_path": "failure-logs/out.txt",
        "stdout_sha256": "4" * 64,
        "stderr_path": "failure-logs/err.txt",
        "stderr_sha256": "5" * 64,
    }
    failure = {
        **failure_body,
        "failure_id": sha256_value(failure_body),
        "recorded_at": "2026-07-20T00:00:00Z",
    }
    command = ["python", "worker.py", "--task-id", recovered["task_id"]]
    recovery_task = {
        "task_id": recovered["task_id"],
        "scientific_run_identity": recovered["scientific_run_identity"],
        "output_relpath": recovered["output_relpath"],
        "execution_config_hash": recovered["execution_config_hash"],
        "original_failure": failure,
        "command": command,
        "recovery_run_key": "sha256:" + "6" * 64,
    }
    recovery_body = {
        "schema_version": 1,
        "kind": "controller_exception_recovery_v1",
        "source_manifest_file_sha256": _file_hash(manifest_path),
        "source_manifest_hash": manifest_hash,
        "source_commit": "7" * 40,
        "recovery_code_commit": "8" * 40,
        "recovery_output_root": outputs.resolve().as_posix(),
        "validator_source_hashes": {
            "artifacts.py": "9" * 64,
            "batch_control.py": "a" * 64,
            "batch_recovery.py": "b" * 64,
        },
        "allowed_failure_reason": failure["reason"],
        "task_count": 1,
        "tasks": [recovery_task],
        "authority": DENIED_AUTHORITY,
    }
    recovery = {
        **recovery_body,
        "recovery_manifest_hash": sha256_value(recovery_body),
    }
    recovery_path = root / "recovery-manifest.json"
    _write_json(recovery_path, recovery)

    recovered_result_path = outputs / recovered["output_relpath"] / "result.json"
    _write_json(recovered_result_path, rows[0])
    result_hash = _file_hash(recovered_result_path)
    artifact_hashes = {"result": result_hash}
    receipt_row = {
        "task_id": recovered["task_id"],
        "recovery_run_key": recovery_task["recovery_run_key"],
        "result_sha256": result_hash,
        "artifact_hashes": artifact_hashes,
        "command": command,
        "returncode": 0,
        "recovery_manifest_hash": recovery["recovery_manifest_hash"],
    }
    success_path = (
        outputs
        / "control"
        / "success"
        / f"{recovery_task['recovery_run_key'].removeprefix('sha256:')}.json"
    )
    _write_json(success_path, receipt_row)
    receipt_row = {**receipt_row, "success_receipt_sha256": _file_hash(success_path)}
    execution_body = {
        "schema_version": 1,
        "kind": "recovery_execution_receipt_v1",
        "recovery_manifest_hash": recovery["recovery_manifest_hash"],
        "task_count": 1,
        "tasks": [receipt_row],
        "authority": DENIED_AUTHORITY,
    }
    execution = {**execution_body, "receipt_hash": sha256_value(execution_body)}
    _write_json(outputs / "recovery-execution-receipt.json", execution)

    attested_task = {
        "task_id": recovered["task_id"],
        "scientific_run_identity": recovered["scientific_run_identity"],
        "original_failure": failure,
        "original_result_sha256": "c" * 64,
        "recovery_result_sha256": result_hash,
        "original_artifact_hashes": {"result": "d" * 64},
        "recovery_artifact_hashes": artifact_hashes,
        "checkpoint_state_sha256": "sha256:" + "e" * 64,
        "scientific_trace_sha256": "sha256:" + "f" * 64,
        "success_receipt_sha256": receipt_row["success_receipt_sha256"],
    }
    attestation_body = {
        "schema_version": 1,
        "kind": "recovery_attestation_v1",
        "source_manifest_hash": manifest_hash,
        "recovery_manifest_hash": recovery["recovery_manifest_hash"],
        "source_commit": recovery["source_commit"],
        "recovery_code_commit": recovery["recovery_code_commit"],
        "reconciliation_code_commit": recovery["recovery_code_commit"],
        "validator": {
            "entrypoint": "loss_aware_dro_repro.batch_control.validate_success_output",
            "source_hashes": recovery["validator_source_hashes"],
        },
        "control_ledger_hashes": {
            "launch-ledger.jsonl": "0" * 64,
            "failure-ledger.jsonl": "1" * 64,
        },
        "tasks": [attested_task],
    }
    attestation = {
        **attestation_body,
        "attestation_hash": sha256_value(attestation_body),
    }
    selected_rows = rows if rows_override is None else rows_override
    result_bytes = b"".join(canonical_bytes(row) + b"\n" for row in selected_rows)
    results_hash = hashlib.sha256(result_bytes).hexdigest()
    aggregate_identity = sha256_value(
        {
            "source_manifest_hash": manifest_hash,
            "results_sha256": results_hash,
            "recovery_attestation_hash": attestation["attestation_hash"],
        }
    )
    aggregate_root = root / "aggregates" / aggregate_identity.removeprefix("sha256:")
    aggregate_root.mkdir(parents=True)
    results_path = aggregate_root / "results.jsonl"
    results_path.write_bytes(result_bytes)
    attestation_path = aggregate_root / "recovery-attestation.json"
    _write_json(attestation_path, attestation)
    summary = {
        "schema_version": 2,
        "kind": "reconciled_batch_aggregate_v1",
        "aggregate_identity": aggregate_identity,
        "source_manifest_hash": manifest_hash,
        "expected_task_count": 2,
        "validated_row_count": 2,
        "recovered_task_count": 1,
        "rejected_row_count": 0,
        "complete": True,
        "artifact_hashes": {
            "results.jsonl": results_hash,
            "recovery-attestation.json": _file_hash(attestation_path),
        },
        "authority": DENIED_AUTHORITY,
        "scientific_verdicts": {f"A{index}": "HOLD" for index in range(1, 7)},
    }
    _write_json(aggregate_root / "summary.json", summary)
    return manifest_path, root, summary


def _validate(manifest_path: Path, root: Path, monkeypatch: pytest.MonkeyPatch):
    import loss_aware_dro_repro.reconciled_analysis as module

    monkeypatch.setattr(module, "validate_batch_manifest", lambda manifest: None)
    return validate_reconciled_bundle(
        manifest_path,
        root,
        _expected_task_count=2,
        _expected_recovered_count=1,
    )


def test_valid_reconciled_bundle_binds_exact_membership_and_substitution(
    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
    manifest_path, root, summary = _fixture(tmp_path)
    _, observed, rows, provenance = _validate(manifest_path, root, monkeypatch)
    assert observed == summary
    assert len(rows) == 2
    assert provenance["validated_row_count"] == 2
    assert provenance["recovered_task_count"] == 1


def test_altered_result_row_is_rejected(
    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
    manifest_path, root, _ = _fixture(tmp_path)
    results_path = next((root / "aggregates").glob("*/results.jsonl"))
    payload = results_path.read_bytes().replace(b'"relative_oos_improvement":0.1', b'"relative_oos_improvement":0.2', 1)
    results_path.write_bytes(payload)
    with pytest.raises(ContractError, match="hash mismatch"):
        _validate(manifest_path, root, monkeypatch)


def test_forged_or_missing_attestation_is_rejected(
    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
    manifest_path, root, _ = _fixture(tmp_path)
    attestation_path = next((root / "aggregates").glob("*/recovery-attestation.json"))
    forged = json.loads(attestation_path.read_text(encoding="utf-8"))
    forged["attestation_hash"] = "sha256:" + "0" * 64
    _write_json(attestation_path, forged)
    with pytest.raises(ContractError, match="attestation hash mismatch"):
        _validate(manifest_path, root, monkeypatch)
    attestation_path.unlink()
    with pytest.raises(ContractError, match="incomplete"):
        _validate(manifest_path, root, monkeypatch)


def test_duplicate_membership_is_rejected(
    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
    seed_path, seed_root, _ = _fixture(tmp_path / "seed")
    first_row = json.loads(next((seed_root / "aggregates").glob("*/results.jsonl")).read_text().splitlines()[0])
    manifest_path, root, _ = _fixture(tmp_path / "duplicate", rows_override=[first_row, copy.deepcopy(first_row)])
    with pytest.raises(ContractError, match="duplicate task identities"):
        _validate(manifest_path, root, monkeypatch)


def test_wrong_manifest_and_recovered_count_are_rejected(
    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
    manifest_path, root, _ = _fixture(tmp_path)
    wrong = tmp_path / "wrong-manifest.json"
    wrong.write_text('{"manifest_hash":"sha256:' + "f" * 64 + '","tasks":[]}', encoding="utf-8")
    with pytest.raises(ContractError, match="different source manifest"):
        _validate(wrong, root, monkeypatch)
    summary_path = next((root / "aggregates").glob("*/summary.json"))
    summary = json.loads(summary_path.read_text(encoding="utf-8"))
    summary["recovered_task_count"] = 2
    _write_json(summary_path, summary)
    with pytest.raises(ContractError, match="wrong task or recovery count"):
        _validate(manifest_path, root, monkeypatch)


def test_declared_recovery_path_escape_is_rejected(
    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
    manifest_path, root, _ = _fixture(tmp_path)
    recovery_path = root / "recovery-manifest.json"
    recovery = json.loads(recovery_path.read_text(encoding="utf-8"))
    recovery["recovery_output_root"] = (tmp_path / "outside").resolve().as_posix()
    body = {key: value for key, value in recovery.items() if key != "recovery_manifest_hash"}
    recovery["recovery_manifest_hash"] = sha256_value(body)
    _write_json(recovery_path, recovery)
    with pytest.raises(ContractError, match="escapes or differs"):
        _validate(manifest_path, root, monkeypatch)