mirror of
https://github.com/HKUDS/nanobot.git
synced 2026-08-31 08:13:11 +03:00
258 lines
8.1 KiB
Python
258 lines
8.1 KiB
Python
from __future__ import annotations
|
|
|
|
import base64
|
|
import json
|
|
import threading
|
|
import time
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
from pathlib import Path
|
|
|
|
import httpx
|
|
import pytest
|
|
|
|
from nanobot.providers.oauth_model_catalog import (
|
|
DEFAULT_XAI_GROK_MODELS_URL,
|
|
OAuthModelCatalog,
|
|
OAuthModelInfo,
|
|
get_oauth_model_catalog,
|
|
invalidate_oauth_model_catalog,
|
|
)
|
|
from nanobot.providers.xai_oauth import XAIToken
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def _clear_xai_catalog() -> None:
|
|
invalidate_oauth_model_catalog("xai_grok")
|
|
yield
|
|
invalidate_oauth_model_catalog("xai_grok")
|
|
|
|
|
|
def _fallback_model() -> OAuthModelInfo:
|
|
return OAuthModelInfo(id="provider/fallback", label="Fallback")
|
|
|
|
|
|
def test_xai_catalog_fetches_remote_models_and_reuses_capability_metadata(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
tmp_path: Path,
|
|
) -> None:
|
|
original_client = httpx.Client
|
|
captured: dict[str, object] = {}
|
|
payload = base64.urlsafe_b64encode(
|
|
json.dumps({"sub": "user-42", "email": "user@example.com"}).encode()
|
|
).decode().rstrip("=")
|
|
token = XAIToken(
|
|
access=f"header.{payload}.signature",
|
|
refresh="refresh-token",
|
|
expires=int(time.time() * 1000) + 3_600_000,
|
|
account_id="user@example.com",
|
|
)
|
|
|
|
def handler(request: httpx.Request) -> httpx.Response:
|
|
captured["request"] = request
|
|
return httpx.Response(
|
|
200,
|
|
json={
|
|
"data": [
|
|
{
|
|
"id": "grok-4.6",
|
|
"name": "Grok 4.6",
|
|
"description": "Latest frontier model",
|
|
"owned_by": "xAI",
|
|
"context_window": 500_000,
|
|
"supports_backend_search": True,
|
|
"reasoning_efforts": [
|
|
{"value": "xhigh"},
|
|
{"value": "high"},
|
|
{"value": "low"},
|
|
],
|
|
},
|
|
{
|
|
"id": "grok-next",
|
|
"_meta": {
|
|
"name": "Grok Next",
|
|
"context_window": 750_000,
|
|
"reasoning_efforts": ["high", "low"],
|
|
},
|
|
},
|
|
]
|
|
},
|
|
request=request,
|
|
)
|
|
|
|
def fake_client(**kwargs: object) -> httpx.Client:
|
|
captured["kwargs"] = kwargs
|
|
return original_client(
|
|
transport=httpx.MockTransport(handler),
|
|
timeout=kwargs["timeout"],
|
|
follow_redirects=kwargs["follow_redirects"],
|
|
)
|
|
|
|
monkeypatch.setattr(
|
|
"nanobot.providers.oauth_model_catalog._xai_oauth_storage_path",
|
|
lambda: tmp_path / "auth" / "xai.json",
|
|
)
|
|
monkeypatch.setattr(
|
|
"nanobot.providers.oauth_model_catalog._xai_oauth_token",
|
|
lambda _proxy: token,
|
|
)
|
|
monkeypatch.setattr("nanobot.providers.oauth_model_catalog.httpx.Client", fake_client)
|
|
|
|
catalog = get_oauth_model_catalog("xai_grok")
|
|
|
|
assert catalog.source == "remote"
|
|
assert [model.id for model in catalog.models] == [
|
|
"xai-grok/grok-4.6",
|
|
"xai-grok/grok-next",
|
|
]
|
|
grok = catalog.find("grok-4.6")
|
|
assert grok is not None
|
|
assert grok.description == "Latest frontier model"
|
|
assert grok.context_window == 500_000
|
|
assert grok.reasoning_efforts == ("xhigh", "high", "low")
|
|
assert grok.supports_backend_search is True
|
|
next_model = catalog.find("xai-grok/grok-next")
|
|
assert next_model is not None
|
|
assert next_model.label == "Grok Next"
|
|
assert next_model.context_window == 750_000
|
|
assert next_model.reasoning_efforts == ("high", "low")
|
|
|
|
request = captured["request"]
|
|
assert isinstance(request, httpx.Request)
|
|
assert str(request.url) == DEFAULT_XAI_GROK_MODELS_URL
|
|
assert request.headers["Authorization"] == f"Bearer {token.access}"
|
|
assert request.headers["X-XAI-Token-Auth"] == "xai-grok-cli"
|
|
assert request.headers["x-userid"] == "user-42"
|
|
assert request.headers["x-email"] == "user@example.com"
|
|
assert captured["kwargs"] == {"timeout": 10.0, "follow_redirects": False}
|
|
assert get_oauth_model_catalog("xai_grok").source == "cache"
|
|
|
|
|
|
def test_catalog_single_flights_concurrent_refreshes() -> None:
|
|
calls = 0
|
|
calls_lock = threading.Lock()
|
|
barrier = threading.Barrier(8)
|
|
|
|
def fetch(_proxy: str | None) -> tuple[OAuthModelInfo, ...]:
|
|
nonlocal calls
|
|
with calls_lock:
|
|
calls += 1
|
|
time.sleep(0.05)
|
|
return (OAuthModelInfo(id="provider/remote", label="Remote"),)
|
|
|
|
catalog = OAuthModelCatalog(fallback_models=(_fallback_model(),), fetch=fetch)
|
|
|
|
def get_catalog(_index: int):
|
|
barrier.wait()
|
|
return catalog.get(cache_key="shared")
|
|
|
|
with ThreadPoolExecutor(max_workers=8) as pool:
|
|
results = list(pool.map(get_catalog, range(8)))
|
|
|
|
assert calls == 1
|
|
assert {result.models[0].id for result in results} == {"provider/remote"}
|
|
assert [result.source for result in results].count("remote") == 1
|
|
assert [result.source for result in results].count("cache") == 7
|
|
|
|
|
|
def test_catalog_returns_stale_then_negative_caches_refresh_failure() -> None:
|
|
now = [0.0]
|
|
calls = 0
|
|
|
|
def fetch(_proxy: str | None) -> tuple[OAuthModelInfo, ...]:
|
|
nonlocal calls
|
|
calls += 1
|
|
if calls > 1:
|
|
raise httpx.ConnectError("offline")
|
|
return (OAuthModelInfo(id="provider/remote", label="Remote"),)
|
|
|
|
catalog = OAuthModelCatalog(
|
|
fallback_models=(_fallback_model(),),
|
|
fetch=fetch,
|
|
fresh_ttl_s=10,
|
|
stale_ttl_s=100,
|
|
failure_ttl_s=30,
|
|
monotonic=lambda: now[0],
|
|
wall_clock=lambda: 123.0,
|
|
)
|
|
|
|
assert catalog.get(cache_key="one").source == "remote"
|
|
now[0] = 11
|
|
stale = catalog.get(cache_key="one")
|
|
assert stale.source == "stale"
|
|
assert stale.models[0].id == "provider/remote"
|
|
assert catalog.get(cache_key="one").source == "stale"
|
|
assert calls == 2
|
|
|
|
now[0] = 101
|
|
fallback = catalog.get(cache_key="one")
|
|
assert fallback.source == "fallback"
|
|
assert fallback.models[0].id == "provider/fallback"
|
|
assert calls == 3
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"failure",
|
|
[
|
|
httpx.ConnectError("offline"),
|
|
ValueError("invalid JSON"),
|
|
httpx.HTTPStatusError(
|
|
"unauthorized",
|
|
request=httpx.Request("GET", DEFAULT_XAI_GROK_MODELS_URL),
|
|
response=httpx.Response(401),
|
|
),
|
|
httpx.HTTPStatusError(
|
|
"rate limited",
|
|
request=httpx.Request("GET", DEFAULT_XAI_GROK_MODELS_URL),
|
|
response=httpx.Response(429),
|
|
),
|
|
httpx.HTTPStatusError(
|
|
"upstream failure",
|
|
request=httpx.Request("GET", DEFAULT_XAI_GROK_MODELS_URL),
|
|
response=httpx.Response(503),
|
|
),
|
|
],
|
|
)
|
|
def test_catalog_falls_back_for_remote_failures(failure: Exception) -> None:
|
|
calls = 0
|
|
|
|
def fetch(_proxy: str | None) -> tuple[OAuthModelInfo, ...]:
|
|
nonlocal calls
|
|
calls += 1
|
|
raise failure
|
|
|
|
catalog = OAuthModelCatalog(
|
|
fallback_models=(_fallback_model(),),
|
|
fetch=fetch,
|
|
failure_ttl_s=30,
|
|
)
|
|
|
|
first = catalog.get(cache_key="one")
|
|
second = catalog.get(cache_key="one")
|
|
|
|
assert first.source == "fallback"
|
|
assert second.source == "fallback"
|
|
assert first.models == (_fallback_model(),)
|
|
assert calls == 1
|
|
|
|
|
|
def test_catalog_treats_empty_remote_list_as_failure_and_can_be_invalidated() -> None:
|
|
calls = 0
|
|
|
|
def fetch(_proxy: str | None) -> tuple[OAuthModelInfo, ...]:
|
|
nonlocal calls
|
|
calls += 1
|
|
return () if calls == 1 else (OAuthModelInfo(id="provider/new", label="New"),)
|
|
|
|
catalog = OAuthModelCatalog(
|
|
fallback_models=(_fallback_model(),),
|
|
fetch=fetch,
|
|
failure_ttl_s=30,
|
|
)
|
|
|
|
assert catalog.get(cache_key="one").source == "fallback"
|
|
catalog.invalidate()
|
|
refreshed = catalog.get(cache_key="one")
|
|
assert refreshed.source == "remote"
|
|
assert refreshed.models[0].id == "provider/new"
|
|
assert calls == 2
|