tests/test_trial.py

"""Trial allowance behaviour, including reinstall and exhaustion."""

# Credit: @paparichens

from __future__ import annotations

from pathlib import Path
from typing import Any

import httpx
import pytest
from test_client import FakeHfClient

from affix_huggingface import (
    AdmissionChecks,
    AdmissionProof,
    AffixHuggingFace,
    ModelRef,
    TrialExhausted,
)
from affix_huggingface.errors import ConfigurationError, RegistrationThrottled
from affix_huggingface.transport import AffixTransport
from affix_huggingface.trial import (
    TrialCredentials,
    load_credentials,
    resolve_subject,
    save_credentials,
    subject_digest,
)

PASSING_CHECKS = AdmissionChecks(
    entitled=True,
    data_route_allowed=True,
    controls_satisfied=True,
)


class FakeQuotaServer:
    """Minimal stand-in for the AffixIO trial endpoints."""

    def __init__(self, *, limit: int = 100) -> None:
        self.limit = limit
        self.subject_used: dict[str, int] = {}
        self.installs: dict[str, str] = {}
        self.install_counter = 0
        self.registrations = 0

    def handler(self, request: httpx.Request) -> httpx.Response:
        path = request.url.path
        if path == "/v1/trial/register":
            return self._register(request)
        install_id = request.headers.get("X-Affix-Install-Id")
        subject = self.installs.get(install_id or "")
        if subject is None:
            return httpx.Response(403, json={"error": "install_rejected"})
        used = self.subject_used.get(subject, 0)
        remaining = max(0, self.limit - used)
        headers = {
            "X-Affix-Trial-Limit": str(self.limit),
            "X-Affix-Trial-Used": str(used),
            "X-Affix-Trial-Remaining": str(remaining),
        }
        if path.endswith("/prove"):
            if remaining <= 0:
                return httpx.Response(
                    403,
                    json={
                        "error": "trial_exhausted",
                        "reason_code": "TRIAL_EXHAUSTED",
                        "message": "Trial allowance used up",
                        "limit": self.limit,
                        "used": used,
                    },
                    headers=headers,
                )
            self.subject_used[subject] = used + 1
            return httpx.Response(200, json={"proof": "proof_hex"}, headers=headers)
        if path == "/v1/gate/verify":
            return httpx.Response(
                200,
                json={"allow": True, "reason_code": "ADMITTED", "mode": "consume"},
                headers=headers,
            )
        if path == "/api/attest":
            return httpx.Response(200, json={"attestation": {"algorithm": "ML-DSA-65"}})
        if path == "/v1/trial/quota":
            return httpx.Response(
                200,
                json={"limit": self.limit, "used": used, "remaining": remaining},
                headers=headers,
            )
        return httpx.Response(200, json={"ok": True}, headers=headers)

    def _register(self, request: httpx.Request) -> httpx.Response:
        import json as jsonlib

        body = jsonlib.loads(request.content.decode("utf-8"))
        subject = str(body["subject"])
        self.registrations += 1
        self.install_counter += 1
        install_id = f"ins_{self.install_counter}"
        self.installs[install_id] = subject
        used = self.subject_used.setdefault(subject, 0)
        return httpx.Response(
            201,
            json={
                "install_id": install_id,
                "install_secret": "secret-value",
                "subject_id": subject,
                "product": "huggingface",
                "quota": {
                    "limit": self.limit,
                    "used": used,
                    "remaining": max(0, self.limit - used),
                },
            },
        )


@pytest.fixture()
def cache_file(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> Path:
    path = tmp_path / "install.json"
    monkeypatch.setenv("AFFIX_INSTALL_FILE", str(path))
    monkeypatch.setenv("AFFIX_TRIAL_SUBJECT", "tester@example.com")
    monkeypatch.setenv("HF_TOKEN", "hf_test_token")
    return path


def build_client(server: FakeQuotaServer) -> tuple[AffixHuggingFace, FakeHfClient]:
    http = httpx.Client(transport=httpx.MockTransport(server.handler))
    transport = AffixTransport("aio_public_trial", client=http)
    hf = FakeHfClient()
    client = AffixHuggingFace(
        affix_transport=transport,
        hf_client=hf,
        trial=True,
        evidence_mode="off",
    )
    return client, hf


def test_first_use_registers_and_caches(cache_file: Path) -> None:
    server = FakeQuotaServer()
    client, _ = build_client(server)

    client.text_generation("prompt", model=ModelRef("org/model"), checks=PASSING_CHECKS)

    assert server.registrations == 1
    assert cache_file.exists()
    cached = load_credentials(cache_file)
    assert cached is not None
    assert cached.install_id == "ins_1"


def test_cached_install_is_reused(cache_file: Path) -> None:
    server = FakeQuotaServer()
    client, _ = build_client(server)

    client.text_generation("one", model=ModelRef("org/model"), checks=PASSING_CHECKS)
    client.text_generation("two", model=ModelRef("org/model"), checks=PASSING_CHECKS)

    assert server.registrations == 1
    assert server.subject_used[subject_digest("tester@example.com")] == 2


def test_deleting_the_cache_does_not_restore_proofs(cache_file: Path) -> None:
    server = FakeQuotaServer(limit=3)
    client, _ = build_client(server)
    for _ in range(3):
        client.text_generation("prompt", model=ModelRef("org/model"), checks=PASSING_CHECKS)

    # Simulate uninstall: local credentials are gone.
    cache_file.unlink()
    assert load_credentials(cache_file) is None

    fresh_client, fresh_hf = build_client(server)
    with pytest.raises(TrialExhausted) as exc:
        fresh_client.text_generation("prompt", model=ModelRef("org/model"), checks=PASSING_CHECKS)

    assert server.registrations == 2
    assert exc.value.limit == 3
    assert fresh_hf.calls == []


def test_exhaustion_blocks_hugging_face(cache_file: Path) -> None:
    server = FakeQuotaServer(limit=2)
    client, hf = build_client(server)
    client.text_generation("one", model=ModelRef("org/model"), checks=PASSING_CHECKS)
    client.text_generation("two", model=ModelRef("org/model"), checks=PASSING_CHECKS)
    assert len(hf.calls) == 2

    with pytest.raises(TrialExhausted):
        client.text_generation("three", model=ModelRef("org/model"), checks=PASSING_CHECKS)
    assert len(hf.calls) == 2


def test_pre_issued_proof_still_reports_quota(cache_file: Path) -> None:
    server = FakeQuotaServer()
    client, _ = build_client(server)
    client.chat_completion(
        [{"role": "user", "content": "hello"}],
        model=ModelRef("org/model"),
        proof=AdmissionProof("proof_hex"),
    )
    quota = client.quota
    assert quota is not None
    assert quota.limit == 100
    assert quota.remaining == 100


def test_refresh_quota_reads_server_state(cache_file: Path) -> None:
    server = FakeQuotaServer(limit=10)
    client, _ = build_client(server)
    client.text_generation("one", model=ModelRef("org/model"), checks=PASSING_CHECKS)

    quota = client.refresh_quota()
    assert quota.used == 1
    assert quota.remaining == 9


def test_install_headers_are_sent(cache_file: Path) -> None:
    seen: list[httpx.Request] = []
    server = FakeQuotaServer()

    def handler(request: httpx.Request) -> httpx.Response:
        seen.append(request)
        return server.handler(request)

    http = httpx.Client(transport=httpx.MockTransport(handler))
    transport = AffixTransport("aio_public_trial", client=http)
    client = AffixHuggingFace(
        affix_transport=transport,
        hf_client=FakeHfClient(),
        trial=True,
        evidence_mode="off",
    )
    client.text_generation("prompt", model=ModelRef("org/model"), checks=PASSING_CHECKS)

    prove = next(r for r in seen if r.url.path.endswith("/prove"))
    assert prove.headers["X-Affix-Install-Id"] == "ins_1"
    assert prove.headers["X-Affix-Install-Secret"] == "secret-value"


def test_registration_throttle_surfaces_clearly(cache_file: Path) -> None:
    def handler(request: httpx.Request) -> httpx.Response:
        return httpx.Response(
            429,
            json={
                "error": "registration_throttled",
                "message": "Too many new trial subjects from this origin today.",
            },
        )

    http = httpx.Client(transport=httpx.MockTransport(handler))
    transport = AffixTransport("aio_public_trial", client=http, max_retries=0)
    client = AffixHuggingFace(
        affix_transport=transport,
        hf_client=FakeHfClient(),
        trial=True,
        evidence_mode="off",
    )
    with pytest.raises(RegistrationThrottled):
        client.text_generation("prompt", model=ModelRef("org/model"), checks=PASSING_CHECKS)


def test_subject_resolution_is_stable_and_opaque() -> None:
    first = resolve_subject(subject="Tester@Example.com")
    second = resolve_subject(subject="  tester@example.com  ")
    assert first == second
    assert "tester" not in first
    assert len(first) == 64


def test_subject_resolution_fails_closed(monkeypatch: pytest.MonkeyPatch) -> None:
    monkeypatch.delenv("AFFIX_TRIAL_SUBJECT", raising=False)
    monkeypatch.delenv("HF_TOKEN", raising=False)
    with pytest.raises(ConfigurationError):
        resolve_subject()


def test_credentials_round_trip(tmp_path: Path) -> None:
    path = tmp_path / "install.json"
    credentials = TrialCredentials(
        install_id="ins_x",
        install_secret="s",
        subject_id="d" * 64,
    )
    save_credentials(credentials, path)
    assert load_credentials(path) == credentials
    assert path.stat().st_mode & 0o777 == 0o600


def test_own_api_key_disables_trial(monkeypatch: pytest.MonkeyPatch) -> None:
    monkeypatch.setenv("HF_TOKEN", "hf_test_token")
    server = FakeQuotaServer()
    http = httpx.Client(transport=httpx.MockTransport(server.handler))
    transport = AffixTransport("aio_customer_key", client=http)
    client = AffixHuggingFace(
        affix_api_key="aio_customer_key",
        affix_transport=transport,
        hf_client=FakeHfClient(),
        evidence_mode="off",
    )
    assert client._trial is False


def test_public_trial_key_used_when_none_supplied(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    monkeypatch.delenv("AFFIX_API_KEY", raising=False)
    monkeypatch.setenv("HF_TOKEN", "hf_test_token")
    captured: dict[str, Any] = {}

    def handler(request: httpx.Request) -> httpx.Response:
        captured["auth"] = request.headers.get("X-API-Key")
        return httpx.Response(200, json={"ok": True})

    http = httpx.Client(transport=httpx.MockTransport(handler))
    client = AffixHuggingFace(hf_client=FakeHfClient(), trial=False)
    client._affix.close()
    client._affix = AffixTransport(client._affix.api_key, client=http)
    client.health()
    assert captured["auth"].startswith("aio_")