nesquena-hermes
83a479d7a5
Codex gate caught that the new load_pool() path runs under the process-default profile (sibling endpoints /api/providers + /api/models/live already wrap). Without it a multi-profile client sees+seeds the default profile's pool (#4247/#4067 profile-isolation class).
400 lines
17 KiB
Python
400 lines
17 KiB
Python
"""Unit tests for PR #4360: Extend credential pool quota status to non-Codex providers.
|
|
|
|
Tests cover:
|
|
- _local_pool_snapshot(provider) function with various pool states
|
|
- Integration with get_provider_quota for non-allowlist providers
|
|
- Preservation of allowlist and openrouter paths
|
|
"""
|
|
import sys
|
|
import types
|
|
import unittest
|
|
from contextlib import contextmanager
|
|
from datetime import datetime, timedelta, timezone
|
|
from types import SimpleNamespace
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
from api.providers import _local_pool_snapshot, get_provider_quota
|
|
|
|
|
|
def _is_ambient_gh_cli_entry_real(source, label, key_source):
|
|
markers = frozenset({"gh_cli", "gh auth token"})
|
|
env_sources = frozenset({"env:github_token", "env:gh_token"})
|
|
source_lower = source.strip().lower()
|
|
return (
|
|
source_lower in markers
|
|
or source_lower in env_sources
|
|
or label.strip().lower() == "gh auth token"
|
|
or key_source.strip().lower() == "gh auth token"
|
|
)
|
|
|
|
|
|
@contextmanager
|
|
def _fake_credential_pool(**kwargs):
|
|
"""Context manager that ensures agent.credential_pool and api.config exist and patches load_pool."""
|
|
had_agent = "agent" in sys.modules
|
|
had_cp = "agent.credential_pool" in sys.modules
|
|
had_api = "api" in sys.modules
|
|
had_api_config = "api.config" in sys.modules
|
|
old_agent = sys.modules.get("agent")
|
|
|
|
if not had_agent:
|
|
_agent = types.ModuleType("agent")
|
|
_agent.__path__ = []
|
|
sys.modules["agent"] = _agent
|
|
if not had_cp:
|
|
_cp = types.ModuleType("agent.credential_pool")
|
|
sys.modules["agent.credential_pool"] = _cp
|
|
sys.modules["agent"].credential_pool = _cp
|
|
|
|
if not had_api:
|
|
_api = types.ModuleType("api")
|
|
_api.__path__ = []
|
|
sys.modules["api"] = _api
|
|
if not had_api_config:
|
|
_api_config = types.ModuleType("api.config")
|
|
_api_config._is_ambient_gh_cli_entry = _is_ambient_gh_cli_entry_real
|
|
sys.modules["api.config"] = _api_config
|
|
sys.modules["api"].config = _api_config
|
|
elif not hasattr(sys.modules["api.config"], "_is_ambient_gh_cli_entry"):
|
|
sys.modules["api.config"]._is_ambient_gh_cli_entry = _is_ambient_gh_cli_entry_real
|
|
|
|
cp = sys.modules["agent.credential_pool"]
|
|
if not hasattr(cp, "load_pool"):
|
|
cp.load_pool = None
|
|
|
|
with patch("agent.credential_pool.load_pool", **kwargs) as m:
|
|
yield m
|
|
|
|
if not had_cp:
|
|
sys.modules.pop("agent.credential_pool", None)
|
|
if had_agent and old_agent is not None and hasattr(old_agent, "credential_pool"):
|
|
delattr(old_agent, "credential_pool")
|
|
if not had_agent:
|
|
sys.modules.pop("agent", None)
|
|
if not had_api_config:
|
|
sys.modules.pop("api.config", None)
|
|
if not had_api:
|
|
sys.modules.pop("api", None)
|
|
|
|
|
|
class TestLocalPoolSnapshot(unittest.TestCase):
|
|
"""Tests for _local_pool_snapshot(provider) function."""
|
|
|
|
def test_local_pool_snapshot_returns_none_for_empty_pool(self):
|
|
"""_local_pool_snapshot returns None when load_pool returns None."""
|
|
with _fake_credential_pool( return_value=None):
|
|
result = _local_pool_snapshot("some-provider")
|
|
assert result is None
|
|
|
|
def test_local_pool_snapshot_returns_none_for_pool_with_no_entries(self):
|
|
"""_local_pool_snapshot returns None when pool has no entries."""
|
|
pool = MagicMock()
|
|
pool.entries.return_value = []
|
|
with _fake_credential_pool( return_value=pool):
|
|
result = _local_pool_snapshot("some-provider")
|
|
assert result is None
|
|
|
|
def test_local_pool_snapshot_with_available_entries(self):
|
|
"""_local_pool_snapshot returns correct snapshot with available entries."""
|
|
entry1 = SimpleNamespace(
|
|
label="Cred 1",
|
|
source="test",
|
|
last_status="available",
|
|
last_status_at=None,
|
|
last_error_code=None,
|
|
last_error_reset_at=None,
|
|
)
|
|
entry2 = SimpleNamespace(
|
|
label="Cred 2",
|
|
source="test",
|
|
last_status="available",
|
|
last_status_at=None,
|
|
last_error_code=None,
|
|
last_error_reset_at=None,
|
|
)
|
|
pool = MagicMock()
|
|
pool.entries.return_value = [entry1, entry2]
|
|
|
|
with _fake_credential_pool( return_value=pool):
|
|
result = _local_pool_snapshot("google")
|
|
assert result is not None
|
|
assert result.provider == "google"
|
|
assert result.source == "local_pool"
|
|
assert result.available is True
|
|
assert result.pool["total_credentials"] == 2
|
|
assert result.pool["available_credentials"] == 2
|
|
assert result.pool["exhausted_credentials"] == 0
|
|
assert len(result.pool["credentials"]) == 2
|
|
assert result.pool["credentials"][0]["status"] == "available"
|
|
assert "2/2 credentials available" in result.details
|
|
|
|
def test_local_pool_snapshot_with_exhausted_entries(self):
|
|
"""_local_pool_snapshot returns correct snapshot with all exhausted entries."""
|
|
exhausted_until = datetime.now(timezone.utc) + timedelta(hours=1)
|
|
entry1 = SimpleNamespace(
|
|
label="Exhausted Cred",
|
|
source="test",
|
|
last_status="exhausted",
|
|
last_status_at=exhausted_until - timedelta(hours=1),
|
|
last_error_code="401",
|
|
last_error_reset_at=exhausted_until,
|
|
)
|
|
pool = MagicMock()
|
|
pool.entries.return_value = [entry1]
|
|
|
|
with _fake_credential_pool( return_value=pool):
|
|
result = _local_pool_snapshot("openai")
|
|
assert result is not None
|
|
assert result.available is False
|
|
assert result.pool["total_credentials"] == 1
|
|
assert result.pool["available_credentials"] == 0
|
|
assert result.pool["exhausted_credentials"] == 1
|
|
assert result.pool["credentials"][0]["status"] == "exhausted"
|
|
assert "All pool credentials are unavailable" in result.unavailable_reason
|
|
assert "1 exhausted" in result.details
|
|
|
|
def test_local_pool_snapshot_with_mixed_entries(self):
|
|
"""_local_pool_snapshot returns correct counts with mixed available/exhausted."""
|
|
entry_avail = SimpleNamespace(
|
|
label="Available",
|
|
source="test",
|
|
last_status="available",
|
|
last_status_at=None,
|
|
last_error_code=None,
|
|
last_error_reset_at=None,
|
|
)
|
|
exhausted_until = datetime.now(timezone.utc) + timedelta(hours=1)
|
|
entry_exhausted = SimpleNamespace(
|
|
label="Exhausted",
|
|
source="test",
|
|
last_status="exhausted",
|
|
last_status_at=exhausted_until - timedelta(hours=1),
|
|
last_error_code="429",
|
|
last_error_reset_at=exhausted_until,
|
|
)
|
|
pool = MagicMock()
|
|
pool.entries.return_value = [entry_avail, entry_exhausted]
|
|
|
|
with _fake_credential_pool( return_value=pool):
|
|
result = _local_pool_snapshot("vertex")
|
|
assert result is not None
|
|
assert result.available is True
|
|
assert result.pool["total_credentials"] == 2
|
|
assert result.pool["available_credentials"] == 1
|
|
assert result.pool["exhausted_credentials"] == 1
|
|
assert "1/2 credentials available" in result.details
|
|
assert "1 exhausted" in result.details
|
|
|
|
def test_local_pool_snapshot_with_dead_entry(self):
|
|
"""_local_pool_snapshot treats dead credentials as unavailable, not available."""
|
|
entry = SimpleNamespace(
|
|
label="Dead Cred",
|
|
source="test",
|
|
last_status="dead",
|
|
last_status_at=None,
|
|
last_error_code="401",
|
|
last_error_reset_at=None,
|
|
)
|
|
pool = MagicMock()
|
|
pool.entries.return_value = [entry]
|
|
|
|
with _fake_credential_pool( return_value=pool):
|
|
result = _local_pool_snapshot("xai-oauth")
|
|
assert result is not None
|
|
assert result.available is False
|
|
assert result.pool["available_credentials"] == 0
|
|
assert result.pool["dead_credentials"] == 1
|
|
assert result.pool["credentials"][0]["status"] == "dead"
|
|
assert "All pool credentials are unavailable" in result.unavailable_reason
|
|
|
|
def test_local_pool_snapshot_handles_load_pool_exception(self):
|
|
"""_local_pool_snapshot returns None when load_pool raises exception."""
|
|
with _fake_credential_pool( side_effect=RuntimeError("Pool error")):
|
|
result = _local_pool_snapshot("custom-provider")
|
|
assert result is None
|
|
|
|
def test_local_pool_snapshot_snapshot_shape(self):
|
|
"""_local_pool_snapshot returns properly shaped SimpleNamespace."""
|
|
entry = SimpleNamespace(
|
|
label="Test",
|
|
source="test",
|
|
last_status="available",
|
|
last_status_at=None,
|
|
last_error_code=None,
|
|
last_error_reset_at=None,
|
|
)
|
|
pool = MagicMock()
|
|
pool.entries.return_value = [entry]
|
|
|
|
with _fake_credential_pool( return_value=pool):
|
|
result = _local_pool_snapshot("test-provider")
|
|
assert hasattr(result, "provider")
|
|
assert hasattr(result, "source")
|
|
assert hasattr(result, "title")
|
|
assert hasattr(result, "plan")
|
|
assert hasattr(result, "windows")
|
|
assert hasattr(result, "details")
|
|
assert hasattr(result, "available")
|
|
assert hasattr(result, "unavailable_reason")
|
|
assert hasattr(result, "fetched_at")
|
|
assert hasattr(result, "pool")
|
|
assert result.title == "Credential pool"
|
|
assert result.plan is None
|
|
assert result.windows == ()
|
|
assert isinstance(result.fetched_at, datetime)
|
|
|
|
|
|
def test_local_pool_snapshot_filters_ambient_gh_cli_entries(self):
|
|
"""_local_pool_snapshot skips ambient gh-cli entries so providers like
|
|
copilot don't show a phantom "1 credential available" (#4247 class)."""
|
|
entry = SimpleNamespace(
|
|
label="gh auth token",
|
|
source="gh_cli",
|
|
key_source="gh auth token",
|
|
last_status="available",
|
|
last_status_at=None,
|
|
last_error_code=None,
|
|
last_error_reset_at=None,
|
|
)
|
|
pool = MagicMock()
|
|
pool.entries.return_value = [entry]
|
|
|
|
with _fake_credential_pool(return_value=pool):
|
|
result = _local_pool_snapshot("copilot")
|
|
assert result is None
|
|
|
|
def test_local_pool_snapshot_keeps_non_ambient_alongside_ambient(self):
|
|
"""Non-ambient entries are kept even when ambient ones are present."""
|
|
ambient = SimpleNamespace(
|
|
label="gh auth token",
|
|
source="gh_cli",
|
|
key_source="gh auth token",
|
|
last_status="available",
|
|
last_status_at=None,
|
|
last_error_code=None,
|
|
last_error_reset_at=None,
|
|
)
|
|
real = SimpleNamespace(
|
|
label="My API Key",
|
|
source="user",
|
|
key_source="",
|
|
last_status="available",
|
|
last_status_at=None,
|
|
last_error_code=None,
|
|
last_error_reset_at=None,
|
|
)
|
|
pool = MagicMock()
|
|
pool.entries.return_value = [ambient, real]
|
|
|
|
with _fake_credential_pool(return_value=pool):
|
|
result = _local_pool_snapshot("copilot")
|
|
assert result is not None
|
|
assert result.pool["total_credentials"] == 1
|
|
assert result.pool["available_credentials"] == 1
|
|
assert len(result.pool["credentials"]) == 1
|
|
|
|
|
|
class TestGetProviderQuotaLocalPool(unittest.TestCase):
|
|
"""Tests for get_provider_quota integration with local pool snapshots."""
|
|
|
|
def test_get_provider_quota_uses_local_pool_for_non_allowlist_provider(self):
|
|
"""get_provider_quota returns available status for non-allowlist provider with pool."""
|
|
entry = SimpleNamespace(
|
|
label="Test Cred",
|
|
source="test",
|
|
last_status="available",
|
|
last_status_at=None,
|
|
last_error_code=None,
|
|
last_error_reset_at=None,
|
|
)
|
|
pool = MagicMock()
|
|
pool.entries.return_value = [entry]
|
|
|
|
with _fake_credential_pool( return_value=pool):
|
|
with patch("api.providers._active_provider_id", return_value="google"):
|
|
result = get_provider_quota("google")
|
|
assert result["ok"] is True
|
|
assert result["provider"] == "google"
|
|
assert result["status"] == "available"
|
|
assert result["supported"] is True
|
|
assert "account_limits" in result
|
|
assert result["account_limits"]["source"] == "local_pool"
|
|
assert result["account_limits"]["available"] is True
|
|
|
|
def test_get_provider_quota_unavailable_when_all_pool_exhausted(self):
|
|
"""get_provider_quota returns unavailable when all pool credentials exhausted."""
|
|
exhausted_until = datetime.now(timezone.utc) + timedelta(hours=1)
|
|
entry = SimpleNamespace(
|
|
label="Exhausted",
|
|
source="test",
|
|
last_status="exhausted",
|
|
last_status_at=exhausted_until - timedelta(hours=1),
|
|
last_error_code="401",
|
|
last_error_reset_at=exhausted_until,
|
|
)
|
|
pool = MagicMock()
|
|
pool.entries.return_value = [entry]
|
|
|
|
with _fake_credential_pool( return_value=pool):
|
|
with patch("api.providers._active_provider_id", return_value="openai"):
|
|
result = get_provider_quota("openai")
|
|
assert result["ok"] is False
|
|
assert result["provider"] == "openai"
|
|
assert result["status"] == "unavailable"
|
|
assert result["supported"] is True
|
|
assert "all credentials are unavailable" in result["message"]
|
|
assert result["account_limits"]["available"] is False
|
|
|
|
def test_get_provider_quota_falls_through_to_unsupported_when_no_pool(self):
|
|
"""get_provider_quota returns unsupported when no pool exists for provider."""
|
|
with _fake_credential_pool( return_value=None):
|
|
with patch("api.providers._active_provider_id", return_value="some-other-provider"):
|
|
result = get_provider_quota("some-other-provider")
|
|
assert result["ok"] is False
|
|
assert result["status"] == "unsupported"
|
|
assert result["supported"] is False
|
|
|
|
def test_get_provider_quota_falls_through_when_pool_has_no_entries(self):
|
|
"""get_provider_quota returns unsupported when pool exists but has no entries."""
|
|
pool = MagicMock()
|
|
pool.entries.return_value = []
|
|
with _fake_credential_pool( return_value=pool):
|
|
with patch("api.providers._active_provider_id", return_value="custom-provider"):
|
|
result = get_provider_quota("custom-provider")
|
|
assert result["ok"] is False
|
|
assert result["status"] == "unsupported"
|
|
|
|
def test_allowlist_providers_bypass_local_pool(self):
|
|
"""get_provider_quota uses probe path for allowlist providers, not local pool."""
|
|
with patch("api.providers._provider_account_usage_status") as mock_probe:
|
|
with _fake_credential_pool(return_value=None) as mock_load_pool:
|
|
mock_probe.return_value = {
|
|
"ok": True,
|
|
"provider": "openai-codex",
|
|
"status": "available",
|
|
}
|
|
get_provider_quota("openai-codex")
|
|
mock_probe.assert_called_once()
|
|
mock_load_pool.assert_not_called()
|
|
|
|
def test_allowlist_anthropic_uses_probe_path(self):
|
|
"""get_provider_quota uses probe path for anthropic, not local pool."""
|
|
with patch("api.providers._provider_account_usage_status") as mock_probe:
|
|
with _fake_credential_pool(return_value=None) as mock_load_pool:
|
|
mock_probe.return_value = {
|
|
"ok": True,
|
|
"provider": "anthropic",
|
|
"status": "available",
|
|
}
|
|
get_provider_quota("anthropic")
|
|
mock_probe.assert_called_once()
|
|
mock_load_pool.assert_not_called()
|
|
|
|
def test_openrouter_uses_api_key_path_not_local_pool(self):
|
|
"""get_provider_quota uses openrouter API-key path, not local pool."""
|
|
with _fake_credential_pool(return_value=None) as mock_load_pool:
|
|
with patch("api.providers._get_provider_api_key", return_value=None):
|
|
result = get_provider_quota("openrouter")
|
|
assert result["status"] == "no_key"
|
|
mock_load_pool.assert_not_called()
|