tests/test_client.py

"""Network-free contract tests for the AffixIO Hugging Face SDK."""

# Credit: @paparichens

from __future__ import annotations

from typing import Any

import pytest

from affix_huggingface import (
    AdmissionChecks,
    AdmissionDenied,
    AdmissionProof,
    AffixHuggingFace,
    ModelRef,
    StreamNotComplete,
)


class FakeAffixTransport:
    def __init__(self, *, allow: bool = True) -> None:
        self.allow = allow
        self.prove_calls: list[dict[str, Any]] = []
        self.gate_calls: list[dict[str, Any]] = []
        self.attest_calls: list[dict[str, Any]] = []

    def close(self) -> None:
        return None

    def health(self) -> tuple[dict[str, Any], str | None]:
        return {"ok": True}, "req_health"

    def prove(self, **kwargs: Any) -> tuple[dict[str, Any], str | None]:
        self.prove_calls.append(kwargs)
        return {"proof": "proof_hex", "valid": True}, "req_prove"

    def gate(self, **kwargs: Any) -> tuple[dict[str, Any], str | None]:
        self.gate_calls.append(kwargs)
        return (
            {
                "allow": self.allow,
                "reason_code": "ADMITTED" if self.allow else "INVALID_PROOF",
                "receipt_id": "rcpt_1",
                "proof_digest": "abc123",
                "proof_ref": "1:abc123",
                "policy_id": "gate.agent",
                "policy_version": "2026.08",
                "mode": kwargs["mode"],
                "spent": self.allow and kwargs["mode"] == "consume",
                "merkle_root": "root_1",
                "merkle_leaf_hash": "leaf_1",
                "attestation": {"algorithm": "ML-DSA-65"},
            },
            "req_gate",
        )

    def attest(self, payload: dict[str, Any], **kwargs: Any) -> tuple[dict[str, Any], str | None]:
        self.attest_calls.append({"payload": payload, **kwargs})
        return {"attestation": {"algorithm": "ML-DSA-65", "signature": "sig"}}, "req_attest"


class FakeHfClient:
    def __init__(self) -> None:
        self.calls: list[dict[str, Any]] = []

    def chat_completion(
        self,
        messages: list[dict[str, Any]],
        *,
        model: str,
        **kwargs: Any,
    ) -> dict[str, Any] | list[str]:
        self.calls.append(
            {"method": "chat_completion", "messages": messages, "model": model, **kwargs}
        )
        if kwargs.get("stream"):
            return ["token-1", "token-2"]
        return {"id": "hf_response_1", "choices": []}

    def text_generation(self, prompt: str, *, model: str, **kwargs: Any) -> str | list[str]:
        self.calls.append(
            {"method": "text_generation", "prompt": prompt, "model": model, **kwargs}
        )
        if kwargs.get("stream"):
            return ["first", "second"]
        return "generated"

    def feature_extraction(
        self,
        text: str | list[str],
        *,
        model: str,
        **kwargs: Any,
    ) -> list[list[float]]:
        self.calls.append(
            {"method": "feature_extraction", "text": text, "model": model, **kwargs}
        )
        return [[0.1, 0.2]]


def make_client(
    *,
    affix: FakeAffixTransport | None = None,
    hf: FakeHfClient | None = None,
    evidence_mode: str = "required",
) -> tuple[AffixHuggingFace, FakeAffixTransport, FakeHfClient]:
    affix = affix or FakeAffixTransport()
    hf = hf or FakeHfClient()
    client = AffixHuggingFace(
        affix_transport=affix,  # type: ignore[arg-type]
        hf_client=hf,
        evidence_mode=evidence_mode,  # type: ignore[arg-type]
    )
    return client, affix, hf


def test_model_binding_is_deterministic_and_excludes_content() -> None:
    model = ModelRef(
        model_id="org/model",
        revision="abc",
        provider="hf-inference",
    )
    assert len(model.binding_digest) == 64
    assert model.binding_digest == ModelRef(
        model_id="org/model",
        revision="abc",
        provider="hf-inference",
    ).binding_digest
    assert "prompt" not in model.binding()


def test_denial_happens_before_hugging_face_call() -> None:
    client, _, hf = make_client(affix=FakeAffixTransport(allow=False))
    with pytest.raises(AdmissionDenied) as exc:
        client.chat_completion(
            [{"role": "user", "content": "private prompt"}],
            model=ModelRef("org/model"),
            proof=AdmissionProof("bad-proof"),
        )
    assert exc.value.reason_code == "INVALID_PROOF"
    assert hf.calls == []


def test_authorised_chat_returns_output_and_evidence() -> None:
    client, affix, hf = make_client()
    result = client.chat_completion(
        [{"role": "user", "content": "private prompt"}],
        model=ModelRef("org/model", revision="commit-sha"),
        proof=AdmissionProof("proof_hex"),
        max_tokens=20,
    )

    assert result.output["id"] == "hf_response_1"
    assert result.receipt.gate.allow is True
    assert result.receipt.completion is not None
    assert hf.calls[0]["messages"][0]["content"] == "private prompt"
    assert "private prompt" not in repr(affix.gate_calls)
    assert "private prompt" not in repr(affix.attest_calls)
    assert affix.attest_calls[0]["payload"]["prompt_included"] is False
    assert affix.attest_calls[0]["payload"]["output_included"] is False


def test_checks_use_remote_prove_then_consume() -> None:
    client, affix, _ = make_client(evidence_mode="off")
    client.text_generation(
        "private input",
        model=ModelRef("org/model"),
        checks=AdmissionChecks(
            entitled=True,
            data_route_allowed=True,
            controls_satisfied=True,
        ),
    )
    assert affix.prove_calls[0]["fields"] == {
        "condition_greater1": True,
        "condition_greater2": True,
        "condition_greater3": True,
    }
    assert affix.gate_calls[0]["mode"] == "consume"
    assert "private input" not in repr(affix.prove_calls)
    assert "private input" not in repr(affix.gate_calls)


def test_stream_attests_only_after_consumption() -> None:
    client, affix, _ = make_client()
    stream = client.stream_text_generation(
        "private input",
        model=ModelRef("org/model"),
        proof=AdmissionProof("proof_hex"),
    )
    with pytest.raises(StreamNotComplete):
        _ = stream.final_receipt
    assert affix.attest_calls == []
    assert list(stream) == ["first", "second"]
    assert stream.complete is True
    assert stream.final_receipt.completion is not None
    assert len(affix.attest_calls) == 1


def test_inspect_does_not_call_hugging_face() -> None:
    client, affix, hf = make_client(evidence_mode="off")
    receipt = client.inspect(
        ModelRef("org/model"),
        proof=AdmissionProof("proof_hex"),
    )
    assert receipt.allow is True
    assert receipt.mode == "check"
    assert affix.gate_calls[0]["mode"] == "check"
    assert hf.calls == []