Adam Moussa
270ce93b2a
Adds the billing-seam invoker (claude_agent_sdk subscription-OAuth, deferred import, API/Bedrock paths) and the Claude-backed clarifier callables (ConfidenceAssessor/QuestionGenerator, one call/turn memoized on (thread_id,len,content-hash), fail-safe to 0.0 so garbage never clears the 98% human gate).
255 lines
8.5 KiB
Python
255 lines
8.5 KiB
Python
"""Unit tests for agent_team.invoker (§3.1) — all mocked, no network/SDK.
|
|
|
|
These tests prove the module imports cleanly without ``claude_agent_sdk`` or
|
|
``anthropic`` installed, and exercise each billing path through the injected
|
|
SDK seams (``_query``/``_options_cls``/``_client``) so nothing real is called.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import pytest
|
|
|
|
from agent_team import billing, invoker
|
|
from agent_team.billing import BillingMode, ClaudeResult
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def _restore_invoker():
|
|
"""Restore the module-global billing invoker after each test."""
|
|
original = billing._invoker
|
|
yield
|
|
billing._invoker = original
|
|
|
|
|
|
# --------------------------------------------------------------------------- #
|
|
# Fakes mirroring the Agent SDK / anthropic message shapes
|
|
# --------------------------------------------------------------------------- #
|
|
|
|
|
|
class _FakeTextBlock:
|
|
def __init__(self, text: str) -> None:
|
|
self.text = text
|
|
|
|
|
|
# Class names mirror the real Agent SDK message types — the invoker dispatches
|
|
# on ``type(msg).__name__``, so these MUST be named AssistantMessage /
|
|
# ResultMessage to be recognised.
|
|
class AssistantMessage:
|
|
def __init__(self, text: str) -> None:
|
|
self.content = [_FakeTextBlock(text)]
|
|
|
|
|
|
class ResultMessage:
|
|
def __init__(self, result: str, cost: float = 0.42) -> None:
|
|
self.result = result
|
|
self.total_cost_usd = cost
|
|
self.usage = {"input_tokens": 11, "output_tokens": 7}
|
|
|
|
|
|
def _fake_options(**kwargs):
|
|
"""Stand-in for ClaudeAgentOptions: just record the kwargs."""
|
|
return dict(kwargs)
|
|
|
|
|
|
def _make_fake_query(messages):
|
|
"""Build an async ``query(prompt=..., options=...)`` yielding ``messages``."""
|
|
|
|
async def _query(*, prompt, options):
|
|
for msg in messages:
|
|
yield msg
|
|
|
|
return _query
|
|
|
|
|
|
class _FakeAnthropicUsage:
|
|
def __init__(self) -> None:
|
|
self.input_tokens = 12
|
|
self.output_tokens = 5
|
|
|
|
def model_dump(self) -> dict:
|
|
return {"input_tokens": self.input_tokens, "output_tokens": self.output_tokens}
|
|
|
|
|
|
class _FakeAnthropicMessage:
|
|
def __init__(self, text: str) -> None:
|
|
self.content = [_FakeTextBlock(text)]
|
|
self.usage = _FakeAnthropicUsage()
|
|
|
|
|
|
class _FakeMessages:
|
|
def __init__(self, text: str) -> None:
|
|
self._text = text
|
|
self.calls: list[dict] = []
|
|
|
|
def create(self, **kwargs):
|
|
self.calls.append(kwargs)
|
|
return _FakeAnthropicMessage(self._text)
|
|
|
|
|
|
class _FakeAnthropicClient:
|
|
def __init__(self, text: str) -> None:
|
|
self.messages = _FakeMessages(text)
|
|
|
|
|
|
# --------------------------------------------------------------------------- #
|
|
# Clean import without the SDKs
|
|
# --------------------------------------------------------------------------- #
|
|
|
|
|
|
def test_module_imports_without_sdks() -> None:
|
|
import importlib
|
|
import sys
|
|
|
|
# claude_agent_sdk / anthropic are not installed in this env.
|
|
assert "claude_agent_sdk" not in sys.modules
|
|
mod = importlib.import_module("agent_team.invoker")
|
|
assert hasattr(mod, "subscription_invoker")
|
|
assert hasattr(mod, "api_invoker")
|
|
|
|
|
|
# --------------------------------------------------------------------------- #
|
|
# Subscription path
|
|
# --------------------------------------------------------------------------- #
|
|
|
|
|
|
def test_subscription_invoker_returns_result(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
monkeypatch.setenv("CLAUDE_CODE_OAUTH_TOKEN", "oauth-tok")
|
|
messages = [
|
|
AssistantMessage("partial "),
|
|
ResultMessage("final answer"),
|
|
]
|
|
result = invoker.subscription_invoker(
|
|
"hello",
|
|
mode=BillingMode.SUBSCRIPTION,
|
|
_query=_make_fake_query(messages),
|
|
_options_cls=_fake_options,
|
|
)
|
|
assert isinstance(result, ClaudeResult)
|
|
assert result.text == "final answer"
|
|
assert result.mode is BillingMode.SUBSCRIPTION
|
|
# usage populated from ResultMessage cost + usage dict.
|
|
assert result.usage["total_cost_usd"] == pytest.approx(0.42)
|
|
assert result.usage["input_tokens"] == 11
|
|
# raw carries the collected message stream.
|
|
assert result.raw == messages
|
|
|
|
|
|
def test_subscription_invoker_falls_back_to_assistant_text(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
monkeypatch.setenv("CLAUDE_CODE_OAUTH_TOKEN", "oauth-tok")
|
|
messages = [AssistantMessage("a"), AssistantMessage("b")]
|
|
result = invoker.subscription_invoker(
|
|
"hi",
|
|
mode=BillingMode.SUBSCRIPTION,
|
|
_query=_make_fake_query(messages),
|
|
_options_cls=_fake_options,
|
|
)
|
|
assert result.text == "a\nb"
|
|
|
|
|
|
def test_subscription_invoker_requires_oauth_token(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
monkeypatch.delenv("CLAUDE_CODE_OAUTH_TOKEN", raising=False)
|
|
with pytest.raises(RuntimeError, match="secrev.env"):
|
|
invoker.subscription_invoker(
|
|
"hi",
|
|
mode=BillingMode.SUBSCRIPTION,
|
|
_query=_make_fake_query([]),
|
|
_options_cls=_fake_options,
|
|
)
|
|
|
|
|
|
# --------------------------------------------------------------------------- #
|
|
# API path
|
|
# --------------------------------------------------------------------------- #
|
|
|
|
|
|
def test_api_invoker_returns_result() -> None:
|
|
client = _FakeAnthropicClient("api text")
|
|
result = invoker.api_invoker("ask", mode=BillingMode.API, _client=client)
|
|
assert result.text == "api text"
|
|
assert result.mode is BillingMode.API
|
|
assert result.usage == {"input_tokens": 12, "output_tokens": 5}
|
|
# The model + prompt were threaded into the SDK call.
|
|
assert client.messages.calls[0]["model"] == invoker.API_MODEL
|
|
assert client.messages.calls[0]["messages"] == [{"role": "user", "content": "ask"}]
|
|
|
|
|
|
# --------------------------------------------------------------------------- #
|
|
# Bedrock path
|
|
# --------------------------------------------------------------------------- #
|
|
|
|
|
|
def test_bedrock_raises_not_implemented() -> None:
|
|
with pytest.raises(NotImplementedError, match="BEDROCK"):
|
|
invoker.real_invoker("hi", mode=BillingMode.BEDROCK)
|
|
|
|
|
|
# --------------------------------------------------------------------------- #
|
|
# real_invoker dispatch — SUBSCRIPTION + API branches forward kwargs to the leaf
|
|
# --------------------------------------------------------------------------- #
|
|
|
|
|
|
def test_real_invoker_dispatches_subscription_branch(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
"""``real_invoker`` routes SUBSCRIPTION to ``subscription_invoker`` (kwargs fwd)."""
|
|
monkeypatch.setenv("CLAUDE_CODE_OAUTH_TOKEN", "oauth-tok")
|
|
messages = [AssistantMessage("partial "), ResultMessage("sub final")]
|
|
result = invoker.real_invoker(
|
|
"hello",
|
|
mode=BillingMode.SUBSCRIPTION,
|
|
_query=_make_fake_query(messages),
|
|
_options_cls=_fake_options,
|
|
)
|
|
assert isinstance(result, ClaudeResult)
|
|
assert result.mode is BillingMode.SUBSCRIPTION
|
|
assert result.text == "sub final" # the ResultMessage leaf ran
|
|
assert result.raw == messages
|
|
|
|
|
|
def test_real_invoker_dispatches_api_branch() -> None:
|
|
"""``real_invoker`` routes API to ``api_invoker`` with the injected client."""
|
|
client = _FakeAnthropicClient("api branch text")
|
|
result = invoker.real_invoker("ask", mode=BillingMode.API, _client=client)
|
|
assert result.mode is BillingMode.API
|
|
assert result.text == "api branch text" # the anthropic-client leaf ran
|
|
# The prompt was threaded through to the injected client.
|
|
assert client.messages.calls[0]["messages"] == [{"role": "user", "content": "ask"}]
|
|
|
|
|
|
# --------------------------------------------------------------------------- #
|
|
# Binding into the billing seam
|
|
# --------------------------------------------------------------------------- #
|
|
|
|
|
|
def test_bind_subscription_invoker_sets_billing_invoker() -> None:
|
|
invoker.bind_subscription_invoker()
|
|
assert billing._invoker is invoker.real_invoker
|
|
|
|
|
|
def test_bound_invoker_drives_claude_invoke(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
# Bind a fake invoker so claude_invoke routes through it end to end.
|
|
captured: dict = {}
|
|
|
|
def fake(prompt: str, *, mode: BillingMode, **kw):
|
|
captured["prompt"] = prompt
|
|
captured["mode"] = mode
|
|
return ClaudeResult(text="routed", mode=mode)
|
|
|
|
invoker.bind_invoker(fake)
|
|
result = billing.claude_invoke("q", mode=BillingMode.API)
|
|
assert result.text == "routed"
|
|
assert captured["mode"] is BillingMode.API
|
|
assert captured["prompt"] == "q"
|
|
|
|
|
|
def test_bind_invoker_defaults_to_real_invoker() -> None:
|
|
invoker.bind_invoker()
|
|
assert billing._invoker is invoker.real_invoker
|