EvoScientist-Multi/tests/test_host_metering_extensions.py

from EvoScientist.llm.errors import AgentControlError
from EvoScientist.middleware.model_fallback import _is_non_fallbackable
import uuid

import httpx
import pytest

from EvoScientist.middleware import recoverable_metering
from EvoScientist.middleware.recoverable_metering import _metering_config, _source_type


def test_agent_control_error_is_non_fallbackable():
    error = AgentControlError(
        "INSUFFICIENT_BALANCE",
        "balance unavailable",
        status_code=403,
    )

    assert "platform control error" in (_is_non_fallbackable(error) or "")
    assert error.model_dump()["code"] == "INSUFFICIENT_BALANCE"


def test_recoverable_metering_reads_explicit_evomemory_scope():
    metering = _metering_config(
        {
            "configurable": {
                "ai4sci_metering": {
                    "gateway_url": "http://gateway",
                    "run_id": "run-parent",
                    "envelope_signature": "signed-parent",
                    "source_type": "evomemory_linker",
                }
            }
        }
    )

    assert metering is not None
    assert metering["source_type"] == "evomemory_linker"
    assert _source_type({"metering_scope": "evomemory_subagent_worker"}, []) == (
        "evomemory_subagent_worker"
    )


@pytest.mark.anyio
async def test_recoverable_metering_sends_internal_service_identity(monkeypatch):
    captured = {}

    class Response:
        def raise_for_status(self):
            return None

    class Client:
        async def post(self, url, *, json, headers):
            captured.update(url=url, json=json, headers=headers)
            return Response()

    monkeypatch.setenv("AI4SCI_EVO_RUNTIME_GRANT_SECRET", "runtime-service-secret")
    monkeypatch.delenv("EVOSCIENTIST_BACKEND_SERVICE_TOKEN", raising=False)
    monkeypatch.setattr(recoverable_metering, "_clients", {"http://gateway": Client()})

    await recoverable_metering._post(
        {
            "gateway_url": "http://gateway",
            "run_id": "run-1",
            "envelope_signature": "signature",
        },
        "start",
        {},
    )

    assert captured["headers"] == {
        "X-Ai4Sci-Service-Token": "runtime-service-secret"
    }


@pytest.mark.anyio
async def test_recoverable_metering_reports_gateway_identity_failure(monkeypatch):
    class Client:
        async def post(self, url, *, json, headers):
            return httpx.Response(401, request=httpx.Request("POST", url))

    monkeypatch.setattr(recoverable_metering, "_clients", {"http://gateway": Client()})

    with pytest.raises(AgentControlError) as exc_info:
        await recoverable_metering._post(
            {
                "gateway_url": "http://gateway",
                "run_id": "run-1",
                "envelope_signature": "signature",
            },
            "start",
            {},
        )

    assert exc_info.value.code == "BILLING_UNAVAILABLE"


async def _fake_sleep(_delay: float) -> None:
    return None


@pytest.mark.anyio
async def test_terminal_metering_preserves_billing_unavailable_and_skips_retry_on_401(monkeypatch):
    calls = {"count": 0}

    class Client:
        async def post(self, url, *, json, headers):
            calls["count"] += 1
            return httpx.Response(401, request=httpx.Request("POST", url))

    monkeypatch.setattr(recoverable_metering, "_clients", {"http://gateway": Client()})
    monkeypatch.setattr(recoverable_metering.asyncio, "sleep", _fake_sleep)

    callback = recoverable_metering.RecoverableMeteringCallback(
        {
            "gateway_url": "http://gateway",
            "run_id": "run-1",
            "envelope_signature": "signature",
        }
    )
    callback._attempts.add(uuid.UUID(int=1))

    with pytest.raises(AgentControlError) as exc_info:
        await callback._terminal(uuid.UUID(int=1), "succeeded", None)

    assert exc_info.value.code == "BILLING_UNAVAILABLE"
    assert calls["count"] == 1