OBLITERATUS/tests/test_checkpoint_capabilities.py

"""Producer-neutral checkpoint capability and dependency diagnostics."""

from __future__ import annotations

import pytest

import obliteratus.checkpoint_capabilities as capability_module
from obliteratus.checkpoint_capabilities import (
    AdapterCapability,
    AdapterRegistry,
    ExactDependency,
    registry_from,
)


def _capability(
    adapter_id: str = "example-dcp",
    *,
    checkpoint_format: str = "pytorch_dcp",
    distribution: str = "example-producer",
    version: str = "1.2.3",
) -> AdapterCapability:
    return AdapterCapability(
        adapter_id=adapter_id,
        adapter_version="0.1.0",
        producer=distribution,
        producer_version=version,
        formats=(checkpoint_format,),
        required_extras=("checkpoint-example",),
        required_dependencies=(ExactDependency(distribution, version),),
    )


def test_empty_registry_does_not_guess_an_adapter_or_dependency():
    resolution = AdapterRegistry().resolve(
        "pytorch_dcp",
        version_provider=lambda _name: pytest.fail("package metadata queried"),
    )

    assert resolution.to_dict() == {
        "status": "missing",
        "adapter_id": None,
        "adapter_version": None,
        "capability_digest": None,
        "reason": (
            "No approved exact capability is registered; structural recognition "
            "does not select an adapter or dependency set."
        ),
    }


@pytest.mark.parametrize("observed", [None, "1.2.2"])
def test_missing_or_incompatible_dependency_has_exact_bounded_diagnostic(observed):
    capability = _capability()
    resolution = AdapterRegistry((capability,)).resolve(
        "pytorch_dcp",
        version_provider=lambda name: observed if name == "example-producer" else None,
        project_version="9.8.7",
        observed_producer="example-producer",
        observed_producer_version="1.2.3",
    )

    assert resolution.status == "missing"
    assert resolution.adapter_id == "example-dcp"
    assert resolution.adapter_version == "0.1.0"
    assert resolution.capability_digest == capability.capability_digest
    assert resolution.reason == (
        "source_identity=verified; dependency_status=missing_or_incompatible; "
        "install_extra=obliteratus[checkpoint-example]==9.8.7; "
        "required_versions=example-producer==1.2.3; "
        f"observed_versions=example-producer={observed or '<missing>'}"
    )


def test_present_exact_dependency_matches_without_granting_trust():
    capability = _capability()
    resolution = AdapterRegistry((capability,)).resolve(
        "pytorch_dcp",
        version_provider=lambda _name: "1.2.3",
        observed_producer="example-producer",
        observed_producer_version="1.2.3",
    )

    assert resolution.status == "matched"
    assert resolution.adapter_id == capability.adapter_id
    assert resolution.capability_digest == capability.capability_digest
    assert "trust authorization is still required" in resolution.reason


def test_format_only_candidate_never_becomes_an_exact_match():
    resolution = AdapterRegistry((_capability(),)).resolve(
        "pytorch_dcp",
        version_provider=lambda _name: "1.2.3",
    )

    assert resolution.status == "missing"
    assert resolution.adapter_id == "example-dcp"
    assert resolution.dependency_unavailable is False
    assert resolution.reason.startswith("source_identity=unverified")


def test_observed_source_identity_must_be_complete_and_match_exactly():
    registry = AdapterRegistry((_capability(),))
    with pytest.raises(ValueError, match="provided together"):
        registry.resolve("pytorch_dcp", observed_producer="example-producer")

    resolution = registry.resolve(
        "pytorch_dcp",
        observed_producer="example-producer",
        observed_producer_version="1.2.4",
        version_provider=lambda _name: pytest.fail("package metadata queried"),
    )
    assert resolution.status == "missing"
    assert resolution.adapter_id is None
    assert "No exact capability matches" in resolution.reason


def test_resolution_and_digest_are_deterministic_for_explicit_records():
    first = _capability(
        "adapter-b",
        checkpoint_format="deepspeed_zero",
        distribution="deepspeed",
        version="0.16.1",
    )
    second = _capability(
        "adapter-a",
        checkpoint_format="megatron_torch_dist",
        distribution="megatron-core",
        version="0.16.1",
    )
    registry = registry_from([first, second])

    assert registry.capabilities == (second, first)
    assert first.capability_digest == _capability(
        "adapter-b",
        checkpoint_format="deepspeed_zero",
        distribution="deepspeed",
        version="0.16.1",
    ).capability_digest


def test_multiple_format_matches_fail_closed_without_dependency_queries():
    registry = AdapterRegistry((_capability("adapter-a"), _capability("adapter-b")))
    resolution = registry.resolve(
        "pytorch_dcp",
        version_provider=lambda _name: pytest.fail("package metadata queried"),
    )

    assert resolution.status == "ambiguous"
    assert resolution.adapter_id is None
    assert resolution.reason == "Multiple exact capabilities match: adapter-a,adapter-b."


@pytest.mark.parametrize(
    ("factory", "message"),
    [
        (lambda: ExactDependency("bad name", "1.0.0"), "distribution"),
        (lambda: ExactDependency("producer", "unselected"), "version"),
        (
            lambda: AdapterCapability(
                "adapter",
                "1.0.0",
                "producer",
                "1.0.0",
                ("hf_safetensors",),
                ("extra",),
                (ExactDependency("producer", "1.0.0"),),
            ),
            "trust-required",
        ),
        (
            lambda: AdapterRegistry((_capability(), _capability())),
            "identifiers must be unique",
        ),
    ],
)
def test_invalid_or_non_exact_capabilities_are_rejected(factory, message):
    with pytest.raises(ValueError, match=message):
        factory()


def test_untrusted_observed_version_is_redacted_from_diagnostic():
    resolution = AdapterRegistry((_capability(),)).resolve(
        "pytorch_dcp",
        version_provider=lambda _name: "bad\nlocal-path=/secret",
    )

    assert "<invalid>" in resolution.reason
    assert "/secret" not in resolution.reason


def test_registry_and_capability_size_limits_are_bounded():
    dependency = ExactDependency("producer", "1.0.0")
    with pytest.raises(ValueError, match="required_extras exceeds"):
        AdapterCapability(
            "adapter",
            "1.0.0",
            "producer",
            "1.0.0",
            ("pytorch_dcp",),
            tuple(f"extra-{index}" for index in range(9)),
            (dependency,),
        )
    with pytest.raises(ValueError, match="registry exceeds"):
        AdapterRegistry(
            tuple(_capability(f"adapter-{index}") for index in range(17))
        )


def test_capability_and_registry_require_immutable_typed_records():
    dependency = ExactDependency("producer", "1.0.0")
    with pytest.raises(TypeError, match="immutable tuples"):
        AdapterCapability(
            "adapter",
            "1.0.0",
            "producer",
            "1.0.0",
            ["pytorch_dcp"],  # type: ignore[arg-type]
            ("extra",),
            (dependency,),
        )
    with pytest.raises(TypeError, match="ExactDependency"):
        AdapterCapability(
            "adapter",
            "1.0.0",
            "producer",
            "1.0.0",
            ("pytorch_dcp",),
            ("extra",),
            ("not-a-record",),  # type: ignore[arg-type]
        )
    with pytest.raises(TypeError, match="immutable tuple"):
        AdapterRegistry([_capability()])  # type: ignore[arg-type]
    with pytest.raises(TypeError, match="AdapterCapability"):
        AdapterRegistry(("not-a-capability",))  # type: ignore[arg-type]


@pytest.mark.parametrize(
    ("formats", "extras", "dependencies", "message"),
    [
        ((), ("extra",), (ExactDependency("producer", "1.0.0"),), "formats must"),
        (("pytorch_dcp",), (), (ExactDependency("producer", "1.0.0"),), "extras must"),
        (("pytorch_dcp",), ("extra",), (), "dependencies must"),
        (("pytorch_dcp",) * 9, ("extra",), (ExactDependency("producer", "1.0.0"),), "formats exceeds"),
        (
            ("pytorch_dcp",),
            ("extra",),
            tuple(ExactDependency(f"producer-{index}", "1.0.0") for index in range(5)),
            "dependencies exceeds",
        ),
        (
            ("pytorch_dcp", "pytorch_dcp"),
            ("extra",),
            (ExactDependency("producer", "1.0.0"),),
            "formats must be unique",
        ),
        (
            ("pytorch_dcp",),
            ("extra", "extra"),
            (ExactDependency("producer", "1.0.0"),),
            "extras must be unique",
        ),
        (
            ("pytorch_dcp",),
            ("extra",),
            (ExactDependency("producer", "1.0.0"),) * 2,
            "distributions must be unique",
        ),
        (
            ("pytorch_dcp",),
            ("bad extra",),
            (ExactDependency("producer", "1.0.0"),),
            "required_extra",
        ),
    ],
)
def test_capability_collections_are_strict_and_bounded(
    formats,
    extras,
    dependencies,
    message,
):
    with pytest.raises(ValueError, match=message):
        AdapterCapability(
            "adapter",
            "1.0.0",
            "producer",
            "1.0.0",
            formats,
            extras,
            dependencies,
        )


def test_installed_version_probe_returns_metadata_or_absence(monkeypatch):
    monkeypatch.setattr(capability_module.metadata, "version", lambda _name: "1.2.3")
    assert capability_module._installed_version("example") == "1.2.3"

    def missing(_name):
        raise capability_module.metadata.PackageNotFoundError

    monkeypatch.setattr(capability_module.metadata, "version", missing)
    assert capability_module._installed_version("example") is None