119 lines
3.7 KiB
Python
119 lines
3.7 KiB
Python
"""Unit tests for agent_team.billing (§3.1)."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import pytest
|
|
|
|
from agent_team import billing
|
|
from agent_team.billing import (
|
|
BillingMode,
|
|
ClaudeResult,
|
|
claude_invoke,
|
|
resolve_mode,
|
|
set_invoker,
|
|
)
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def _restore_invoker():
|
|
"""Restore the module invoker after each test."""
|
|
original = billing._invoker
|
|
yield
|
|
billing._invoker = original
|
|
|
|
|
|
def test_resolve_mode_default_is_subscription(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
monkeypatch.delenv("AGENT_TEAM_BILLING_MODE", raising=False)
|
|
assert resolve_mode(None) is BillingMode.SUBSCRIPTION
|
|
assert resolve_mode({}) is BillingMode.SUBSCRIPTION
|
|
|
|
|
|
def test_resolve_mode_from_config_string() -> None:
|
|
assert resolve_mode({"billing_mode": "api"}) is BillingMode.API
|
|
assert resolve_mode({"billing_mode": "BEDROCK"}) is BillingMode.BEDROCK
|
|
|
|
|
|
def test_resolve_mode_from_config_enum() -> None:
|
|
assert resolve_mode({"billing_mode": BillingMode.API}) is BillingMode.API
|
|
|
|
|
|
def test_resolve_mode_from_env(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
monkeypatch.setenv("AGENT_TEAM_BILLING_MODE", "bedrock")
|
|
assert resolve_mode(None) is BillingMode.BEDROCK
|
|
|
|
|
|
def test_resolve_mode_config_overrides_env(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
monkeypatch.setenv("AGENT_TEAM_BILLING_MODE", "bedrock")
|
|
assert resolve_mode({"billing_mode": "api"}) is BillingMode.API
|
|
|
|
|
|
def test_resolve_mode_invalid_raises() -> None:
|
|
with pytest.raises(ValueError):
|
|
resolve_mode({"billing_mode": "carrier-pigeon"})
|
|
|
|
|
|
def test_claude_invoke_delegates_with_resolved_mode() -> None:
|
|
captured: dict = {}
|
|
|
|
def fake(prompt: str, *, mode: BillingMode, **kw):
|
|
captured["prompt"] = prompt
|
|
captured["mode"] = mode
|
|
captured["kw"] = kw
|
|
return ClaudeResult(text="ok", mode=mode)
|
|
|
|
set_invoker(fake)
|
|
result = claude_invoke("hi", mode=BillingMode.API, temperature=0.2)
|
|
assert result.text == "ok"
|
|
assert captured["mode"] is BillingMode.API
|
|
assert captured["prompt"] == "hi"
|
|
assert captured["kw"] == {"temperature": 0.2}
|
|
|
|
|
|
def test_subscription_mode_pops_stray_api_key(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
monkeypatch.setenv("ANTHROPIC_API_KEY", "sk-should-be-hidden")
|
|
seen: dict = {}
|
|
|
|
def fake(prompt: str, *, mode: BillingMode, **kw):
|
|
import os
|
|
|
|
seen["key_present"] = "ANTHROPIC_API_KEY" in os.environ
|
|
return ClaudeResult(text="ok", mode=mode)
|
|
|
|
set_invoker(fake)
|
|
claude_invoke("hi", mode=BillingMode.SUBSCRIPTION)
|
|
assert seen["key_present"] is False
|
|
# Restored after the call.
|
|
import os
|
|
|
|
assert os.environ.get("ANTHROPIC_API_KEY") == "sk-should-be-hidden"
|
|
|
|
|
|
def test_api_mode_does_not_pop_api_key(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
monkeypatch.setenv("ANTHROPIC_API_KEY", "sk-metered")
|
|
seen: dict = {}
|
|
|
|
def fake(prompt: str, *, mode: BillingMode, **kw):
|
|
import os
|
|
|
|
seen["key_present"] = "ANTHROPIC_API_KEY" in os.environ
|
|
return ClaudeResult(text="ok", mode=mode)
|
|
|
|
set_invoker(fake)
|
|
claude_invoke("hi", mode=BillingMode.API)
|
|
assert seen["key_present"] is True
|
|
|
|
|
|
def test_subscription_with_no_key_is_safe(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
monkeypatch.delenv("ANTHROPIC_API_KEY", raising=False)
|
|
set_invoker(lambda prompt, *, mode, **kw: ClaudeResult(text="ok", mode=mode))
|
|
assert claude_invoke("hi", mode=BillingMode.SUBSCRIPTION).text == "ok"
|
|
|
|
|
|
def test_unconfigured_invoker_raises() -> None:
|
|
billing._invoker = billing._unconfigured_invoker
|
|
with pytest.raises(RuntimeError):
|
|
claude_invoke("hi", mode=BillingMode.API)
|
|
|
|
|
|
def test_billing_mode_enum_members() -> None:
|
|
assert {m.name for m in BillingMode} == {"SUBSCRIPTION", "API", "BEDROCK"}
|