wehub-resource-sync
636 lines
20 KiB
Python
636 lines
20 KiB
Python
import asyncio
|
|
from typing import cast
|
|
from unittest.mock import AsyncMock, MagicMock, patch
|
|
|
|
import pytest
|
|
|
|
from free_claude_code.application.errors import ApplicationUnavailableError
|
|
from free_claude_code.application.model_metadata import ProviderModelInfo
|
|
from free_claude_code.config.settings import Settings
|
|
from free_claude_code.providers.base import BaseProvider
|
|
from free_claude_code.providers.nvidia_nim import NvidiaNimProvider
|
|
from free_claude_code.providers.runtime import ProviderRuntime
|
|
from free_claude_code.runtime.provider_manager import ProviderRuntimeManager
|
|
|
|
|
|
class FakeRuntime(ProviderRuntime):
|
|
def __init__(self, settings: Settings) -> None:
|
|
self.settings = settings
|
|
self.cleanup_calls = 0
|
|
self.cleanup_error: Exception | None = None
|
|
self.cleanup_started: asyncio.Event | None = None
|
|
self.cleanup_release: asyncio.Event | None = None
|
|
self.provider = MagicMock()
|
|
self.provider.list_model_infos = AsyncMock(return_value=frozenset())
|
|
|
|
def is_cached(self, provider_id: str) -> bool:
|
|
return provider_id == "cached"
|
|
|
|
def resolve_provider(self, provider_id: str) -> BaseProvider:
|
|
return cast(BaseProvider, self.provider)
|
|
|
|
async def cleanup(self) -> None:
|
|
self.cleanup_calls += 1
|
|
if self.cleanup_started is not None:
|
|
self.cleanup_started.set()
|
|
if self.cleanup_release is not None:
|
|
await self.cleanup_release.wait()
|
|
if self.cleanup_error is not None:
|
|
raise self.cleanup_error
|
|
|
|
|
|
class RuntimeFactory:
|
|
def __init__(self) -> None:
|
|
self.runtimes: list[FakeRuntime] = []
|
|
self.error: Exception | None = None
|
|
|
|
def __call__(self, settings: Settings) -> ProviderRuntime:
|
|
if self.error is not None:
|
|
raise self.error
|
|
runtime = FakeRuntime(settings)
|
|
self.runtimes.append(runtime)
|
|
return runtime
|
|
|
|
|
|
def _settings(model: str) -> Settings:
|
|
return Settings().model_copy(update={"model": model})
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_startup_generation_lease_and_shutdown_close_exactly_once() -> None:
|
|
factory = RuntimeFactory()
|
|
settings = _settings("nvidia_nim/one")
|
|
manager = ProviderRuntimeManager(settings, runtime_factory=factory)
|
|
|
|
lease = await manager.acquire()
|
|
|
|
assert lease.generation_id == 1
|
|
assert lease.settings is settings
|
|
assert lease.is_provider_cached("cached") is True
|
|
assert lease.resolve_provider("nvidia_nim") is factory.runtimes[0].provider
|
|
await lease.release()
|
|
await lease.release()
|
|
await manager.close()
|
|
await manager.close()
|
|
|
|
assert factory.runtimes[0].cleanup_calls == 1
|
|
with pytest.raises(ApplicationUnavailableError, match="shutting down"):
|
|
await manager.acquire()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_replacement_keeps_leased_generation_open_until_final_release() -> None:
|
|
factory = RuntimeFactory()
|
|
first_settings = _settings("nvidia_nim/one")
|
|
second_settings = _settings("nvidia_nim/two")
|
|
manager = ProviderRuntimeManager(first_settings, runtime_factory=factory)
|
|
old_lease = await manager.acquire()
|
|
committed: list[str] = []
|
|
|
|
generation_id = await manager.replace(
|
|
second_settings,
|
|
commit=lambda: committed.append("persisted"),
|
|
)
|
|
new_lease = await manager.acquire()
|
|
|
|
assert generation_id == 2
|
|
assert committed == ["persisted"]
|
|
assert new_lease.generation_id == 2
|
|
assert new_lease.settings is second_settings
|
|
assert factory.runtimes[0].cleanup_calls == 0
|
|
await new_lease.release()
|
|
await old_lease.release()
|
|
assert factory.runtimes[0].cleanup_calls == 1
|
|
await manager.close()
|
|
assert factory.runtimes[1].cleanup_calls == 1
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_real_hot_replacement_owns_a_limiter_per_provider_generation() -> None:
|
|
first_settings = _settings("nvidia_nim/one")
|
|
second_settings = _settings("nvidia_nim/two")
|
|
clients: list[MagicMock] = []
|
|
|
|
def create_client(*_args: object, **_kwargs: object) -> MagicMock:
|
|
client = MagicMock()
|
|
client.close = AsyncMock()
|
|
clients.append(client)
|
|
return client
|
|
|
|
with patch(
|
|
"free_claude_code.providers.openai_chat.provider.AsyncOpenAI",
|
|
side_effect=create_client,
|
|
):
|
|
manager = ProviderRuntimeManager(first_settings)
|
|
old_lease = await manager.acquire()
|
|
old_provider = old_lease.resolve_provider("nvidia_nim")
|
|
refresh = AsyncMock()
|
|
|
|
with patch.object(manager, "_refresh_generation", refresh):
|
|
await manager.replace(second_settings, commit=lambda: None)
|
|
new_lease = await manager.acquire()
|
|
new_provider = new_lease.resolve_provider("nvidia_nim")
|
|
await asyncio.sleep(0)
|
|
|
|
assert isinstance(old_provider, NvidiaNimProvider)
|
|
assert isinstance(new_provider, NvidiaNimProvider)
|
|
assert new_provider is not old_provider
|
|
assert new_provider._rate_limiter is not old_provider._rate_limiter
|
|
assert old_lease.resolve_provider("nvidia_nim") is old_provider
|
|
clients[0].close.assert_not_awaited()
|
|
|
|
await new_lease.release()
|
|
await old_lease.release()
|
|
|
|
clients[0].close.assert_awaited_once()
|
|
clients[1].close.assert_not_awaited()
|
|
refresh.assert_awaited_once()
|
|
await manager.close()
|
|
|
|
clients[1].close.assert_awaited_once()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_replacement_closes_unleased_generation_immediately() -> None:
|
|
factory = RuntimeFactory()
|
|
manager = ProviderRuntimeManager(
|
|
_settings("nvidia_nim/one"),
|
|
runtime_factory=factory,
|
|
)
|
|
|
|
await manager.replace(
|
|
_settings("nvidia_nim/two"),
|
|
commit=lambda: None,
|
|
)
|
|
|
|
assert factory.runtimes[0].cleanup_calls == 1
|
|
await manager.close()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_cancelled_replacement_does_not_cancel_owned_generation_cleanup() -> None:
|
|
factory = RuntimeFactory()
|
|
manager = ProviderRuntimeManager(
|
|
_settings("nvidia_nim/one"),
|
|
runtime_factory=factory,
|
|
)
|
|
cleanup_started = asyncio.Event()
|
|
cleanup_release = asyncio.Event()
|
|
refresh_started = asyncio.Event()
|
|
factory.runtimes[0].cleanup_started = cleanup_started
|
|
factory.runtimes[0].cleanup_release = cleanup_release
|
|
|
|
async def refresh(*_args: object, **_kwargs: object) -> None:
|
|
refresh_started.set()
|
|
await asyncio.Event().wait()
|
|
|
|
with patch.object(manager, "_refresh_generation", side_effect=refresh):
|
|
replace_task = asyncio.create_task(
|
|
manager.replace(
|
|
_settings("nvidia_nim/two"),
|
|
commit=lambda: None,
|
|
)
|
|
)
|
|
await cleanup_started.wait()
|
|
await refresh_started.wait()
|
|
|
|
replace_task.cancel()
|
|
with pytest.raises(asyncio.CancelledError):
|
|
await replace_task
|
|
|
|
retired = manager._retired[1]
|
|
assert manager.current_generation_id == 2
|
|
assert retired.cleanup_task is not None
|
|
assert not retired.cleanup_task.cancelled()
|
|
assert factory.runtimes[0].cleanup_calls == 1
|
|
|
|
close_task = asyncio.create_task(manager.close())
|
|
await asyncio.sleep(0)
|
|
assert not close_task.done()
|
|
cleanup_release.set()
|
|
await close_task
|
|
|
|
assert factory.runtimes[0].cleanup_calls == 1
|
|
assert factory.runtimes[1].cleanup_calls == 1
|
|
assert manager._retired == {}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_cancelled_final_lease_release_keeps_owned_cleanup_running() -> None:
|
|
factory = RuntimeFactory()
|
|
manager = ProviderRuntimeManager(
|
|
_settings("nvidia_nim/one"),
|
|
runtime_factory=factory,
|
|
)
|
|
lease = await manager.acquire()
|
|
await manager.replace(
|
|
_settings("nvidia_nim/two"),
|
|
commit=lambda: None,
|
|
)
|
|
cleanup_started = asyncio.Event()
|
|
cleanup_release = asyncio.Event()
|
|
factory.runtimes[0].cleanup_started = cleanup_started
|
|
factory.runtimes[0].cleanup_release = cleanup_release
|
|
|
|
release_task = asyncio.create_task(lease.release())
|
|
await cleanup_started.wait()
|
|
release_task.cancel()
|
|
|
|
with pytest.raises(asyncio.CancelledError):
|
|
await release_task
|
|
|
|
retired = manager._retired[1]
|
|
assert retired.active_leases == 0
|
|
assert retired.cleanup_task is not None
|
|
assert not retired.cleanup_task.cancelled()
|
|
assert factory.runtimes[0].cleanup_calls == 1
|
|
|
|
close_task = asyncio.create_task(manager.close())
|
|
await asyncio.sleep(0)
|
|
assert not close_task.done()
|
|
cleanup_release.set()
|
|
await close_task
|
|
|
|
assert factory.runtimes[0].cleanup_calls == 1
|
|
assert manager._retired == {}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_hot_cleanup_failure_keeps_published_replacement() -> None:
|
|
factory = RuntimeFactory()
|
|
manager = ProviderRuntimeManager(
|
|
_settings("nvidia_nim/one"),
|
|
runtime_factory=factory,
|
|
)
|
|
factory.runtimes[0].cleanup_error = RuntimeError("cleanup failed")
|
|
|
|
generation_id = await manager.replace(
|
|
_settings("nvidia_nim/two"),
|
|
commit=lambda: None,
|
|
)
|
|
|
|
assert generation_id == 2
|
|
assert manager.current_generation_id == 2
|
|
assert factory.runtimes[0].cleanup_calls == 1
|
|
assert 1 in manager._retired
|
|
|
|
factory.runtimes[0].cleanup_error = None
|
|
await manager.close()
|
|
|
|
assert factory.runtimes[0].cleanup_calls == 2
|
|
assert manager._retired == {}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_candidate_construction_failure_preserves_current_generation() -> None:
|
|
factory = RuntimeFactory()
|
|
manager = ProviderRuntimeManager(
|
|
_settings("nvidia_nim/one"),
|
|
runtime_factory=factory,
|
|
)
|
|
factory.error = RuntimeError("cannot construct")
|
|
|
|
with pytest.raises(RuntimeError, match="cannot construct"):
|
|
await manager.replace(
|
|
_settings("nvidia_nim/two"),
|
|
commit=lambda: None,
|
|
)
|
|
|
|
assert manager.current_generation_id == 1
|
|
assert manager.current_settings().model == "nvidia_nim/one"
|
|
assert factory.runtimes[0].cleanup_calls == 0
|
|
await manager.close()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_failed_candidate_cleanup_is_retried_at_shutdown() -> None:
|
|
factory = RuntimeFactory()
|
|
manager = ProviderRuntimeManager(
|
|
_settings("nvidia_nim/one"),
|
|
runtime_factory=factory,
|
|
)
|
|
|
|
def fail_commit() -> None:
|
|
factory.runtimes[1].cleanup_error = RuntimeError("private cleanup detail")
|
|
raise OSError("disk full")
|
|
|
|
with pytest.raises(OSError, match="disk full"):
|
|
await manager.replace(
|
|
_settings("nvidia_nim/two"),
|
|
commit=fail_commit,
|
|
)
|
|
|
|
assert manager.current_generation_id == 1
|
|
assert factory.runtimes[0].cleanup_calls == 0
|
|
assert factory.runtimes[1].cleanup_calls == 1
|
|
assert manager._unpublished == {factory.runtimes[1]}
|
|
|
|
manager.cache_model_infos("nvidia_nim", {ProviderModelInfo("cached")})
|
|
with pytest.raises(
|
|
RuntimeError,
|
|
match="One or more provider runtimes failed to close",
|
|
) as exc_info:
|
|
await manager.close()
|
|
|
|
assert "private cleanup detail" not in str(exc_info.value)
|
|
assert manager._closed is False
|
|
assert manager._unpublished == {factory.runtimes[1]}
|
|
assert manager.cached_model_ids() == {"nvidia_nim": frozenset({"cached"})}
|
|
|
|
factory.runtimes[1].cleanup_error = None
|
|
await manager.close()
|
|
|
|
assert factory.runtimes[0].cleanup_calls == 1
|
|
assert factory.runtimes[1].cleanup_calls == 3
|
|
assert manager._unpublished == set()
|
|
assert manager.cached_model_ids() == {}
|
|
assert manager._closed is True
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_later_replacement_retries_failed_unpublished_candidate() -> None:
|
|
factory = RuntimeFactory()
|
|
manager = ProviderRuntimeManager(
|
|
_settings("nvidia_nim/one"),
|
|
runtime_factory=factory,
|
|
)
|
|
|
|
def fail_commit() -> None:
|
|
factory.runtimes[1].cleanup_error = RuntimeError("cleanup failed")
|
|
raise OSError("disk full")
|
|
|
|
with pytest.raises(OSError, match="disk full"):
|
|
await manager.replace(
|
|
_settings("nvidia_nim/two"),
|
|
commit=fail_commit,
|
|
)
|
|
|
|
factory.runtimes[1].cleanup_error = None
|
|
generation_id = await manager.replace(
|
|
_settings("nvidia_nim/three"),
|
|
commit=lambda: None,
|
|
)
|
|
|
|
assert generation_id == 2
|
|
assert factory.runtimes[1].cleanup_calls == 2
|
|
assert manager._unpublished == set()
|
|
await manager.close()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_cancelled_candidate_cleanup_remains_owned_until_shutdown() -> None:
|
|
factory = RuntimeFactory()
|
|
manager = ProviderRuntimeManager(
|
|
_settings("nvidia_nim/one"),
|
|
runtime_factory=factory,
|
|
)
|
|
cleanup_started = asyncio.Event()
|
|
cleanup_release = asyncio.Event()
|
|
|
|
def fail_commit() -> None:
|
|
candidate = factory.runtimes[1]
|
|
candidate.cleanup_started = cleanup_started
|
|
candidate.cleanup_release = cleanup_release
|
|
raise OSError("disk full")
|
|
|
|
replace_task = asyncio.create_task(
|
|
manager.replace(
|
|
_settings("nvidia_nim/two"),
|
|
commit=fail_commit,
|
|
)
|
|
)
|
|
await cleanup_started.wait()
|
|
replace_task.cancel()
|
|
|
|
with pytest.raises(asyncio.CancelledError):
|
|
await replace_task
|
|
|
|
assert manager._unpublished == {factory.runtimes[1]}
|
|
assert factory.runtimes[1].cleanup_calls == 1
|
|
|
|
cleanup_release.set()
|
|
await manager.close()
|
|
|
|
assert factory.runtimes[1].cleanup_calls == 2
|
|
assert manager._unpublished == set()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_concurrent_replacements_are_serialized_in_call_order() -> None:
|
|
factory = RuntimeFactory()
|
|
manager = ProviderRuntimeManager(
|
|
_settings("nvidia_nim/one"),
|
|
runtime_factory=factory,
|
|
)
|
|
first_entered = asyncio.Event()
|
|
release_first = asyncio.Event()
|
|
cancel_calls = 0
|
|
original_cancel = manager._cancel_refresh
|
|
|
|
async def controlled_cancel() -> None:
|
|
nonlocal cancel_calls
|
|
cancel_calls += 1
|
|
if cancel_calls == 1:
|
|
first_entered.set()
|
|
await release_first.wait()
|
|
await original_cancel()
|
|
|
|
with patch.object(manager, "_cancel_refresh", side_effect=controlled_cancel):
|
|
first = asyncio.create_task(
|
|
manager.replace(
|
|
_settings("nvidia_nim/two"),
|
|
commit=lambda: None,
|
|
)
|
|
)
|
|
await first_entered.wait()
|
|
second = asyncio.create_task(
|
|
manager.replace(
|
|
_settings("nvidia_nim/three"),
|
|
commit=lambda: None,
|
|
)
|
|
)
|
|
await asyncio.sleep(0)
|
|
assert len(factory.runtimes) == 1
|
|
assert not second.done()
|
|
release_first.set()
|
|
assert await asyncio.gather(first, second) == [2, 3]
|
|
|
|
assert manager.current_settings().model == "nvidia_nim/three"
|
|
assert [runtime.cleanup_calls for runtime in factory.runtimes[:2]] == [1, 1]
|
|
await manager.close()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_shutdown_waits_for_active_lease_then_rejects_new_work() -> None:
|
|
factory = RuntimeFactory()
|
|
manager = ProviderRuntimeManager(
|
|
_settings("nvidia_nim/one"),
|
|
runtime_factory=factory,
|
|
)
|
|
lease = await manager.acquire()
|
|
|
|
close_task = asyncio.create_task(manager.close())
|
|
await asyncio.sleep(0)
|
|
|
|
assert not close_task.done()
|
|
with pytest.raises(ApplicationUnavailableError, match="shutting down"):
|
|
await manager.acquire()
|
|
await lease.release()
|
|
await close_task
|
|
assert factory.runtimes[0].cleanup_calls == 1
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_cancelled_shutdown_retains_generation_for_retry() -> None:
|
|
factory = RuntimeFactory()
|
|
manager = ProviderRuntimeManager(
|
|
_settings("nvidia_nim/one"),
|
|
runtime_factory=factory,
|
|
)
|
|
lease = await manager.acquire()
|
|
close_task = asyncio.create_task(manager.close())
|
|
await asyncio.sleep(0)
|
|
|
|
close_task.cancel()
|
|
with pytest.raises(asyncio.CancelledError):
|
|
await close_task
|
|
|
|
with pytest.raises(ApplicationUnavailableError, match="shutting down"):
|
|
await manager.acquire()
|
|
with pytest.raises(ApplicationUnavailableError, match="shutting down"):
|
|
await manager.replace(_settings("nvidia_nim/two"), commit=lambda: None)
|
|
assert factory.runtimes[0].cleanup_calls == 0
|
|
|
|
await lease.release()
|
|
await manager.close()
|
|
|
|
assert factory.runtimes[0].cleanup_calls == 1
|
|
assert manager._closed is True
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_cancelled_shutdown_reuses_the_same_owned_cleanup_task() -> None:
|
|
factory = RuntimeFactory()
|
|
manager = ProviderRuntimeManager(
|
|
_settings("nvidia_nim/one"),
|
|
runtime_factory=factory,
|
|
)
|
|
cleanup_started = asyncio.Event()
|
|
cleanup_release = asyncio.Event()
|
|
cleanup_calls = 0
|
|
|
|
async def cleanup() -> None:
|
|
nonlocal cleanup_calls
|
|
cleanup_calls += 1
|
|
cleanup_started.set()
|
|
await cleanup_release.wait()
|
|
|
|
with patch.object(factory.runtimes[0], "cleanup", side_effect=cleanup):
|
|
close_task = asyncio.create_task(manager.close())
|
|
await cleanup_started.wait()
|
|
|
|
close_task.cancel()
|
|
with pytest.raises(asyncio.CancelledError):
|
|
await close_task
|
|
|
|
assert manager._closed is False
|
|
assert manager._retired
|
|
|
|
retry_task = asyncio.create_task(manager.close())
|
|
await asyncio.sleep(0)
|
|
assert not retry_task.done()
|
|
cleanup_release.set()
|
|
await retry_task
|
|
|
|
assert cleanup_calls == 1
|
|
assert manager._retired == {}
|
|
assert manager._closed is True
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_failed_shutdown_cleanup_is_retryable() -> None:
|
|
factory = RuntimeFactory()
|
|
manager = ProviderRuntimeManager(
|
|
_settings("nvidia_nim/one"),
|
|
runtime_factory=factory,
|
|
)
|
|
manager.cache_model_infos("nvidia_nim", {ProviderModelInfo("cached")})
|
|
factory.runtimes[0].cleanup_error = RuntimeError("private provider detail")
|
|
|
|
with pytest.raises(
|
|
RuntimeError,
|
|
match="One or more provider runtimes failed to close",
|
|
) as exc_info:
|
|
await manager.close()
|
|
|
|
assert "private provider detail" not in str(exc_info.value)
|
|
assert manager._closed is False
|
|
assert 1 in manager._retired
|
|
assert manager.cached_model_ids() == {"nvidia_nim": frozenset({"cached"})}
|
|
assert factory.runtimes[0].cleanup_calls == 1
|
|
|
|
factory.runtimes[0].cleanup_error = None
|
|
await manager.close()
|
|
|
|
assert factory.runtimes[0].cleanup_calls == 2
|
|
assert manager._retired == {}
|
|
assert manager.cached_model_ids() == {}
|
|
assert manager._closed is True
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_application_catalog_survives_generation_replacement() -> None:
|
|
factory = RuntimeFactory()
|
|
manager = ProviderRuntimeManager(
|
|
_settings("open_router/one"),
|
|
runtime_factory=factory,
|
|
)
|
|
manager.cache_model_infos(
|
|
"open_router",
|
|
{ProviderModelInfo("persisted", supports_thinking=True)},
|
|
)
|
|
|
|
await manager.replace(
|
|
_settings("open_router/two"),
|
|
commit=lambda: None,
|
|
)
|
|
|
|
assert manager.cached_model_ids() == {"open_router": frozenset({"persisted"})}
|
|
assert manager.cached_model_supports_thinking("open_router", "persisted") is True
|
|
await manager.close()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_generation_lifecycle_traces_contain_minimal_correlation_fields() -> None:
|
|
factory = RuntimeFactory()
|
|
|
|
with patch("free_claude_code.runtime.provider_manager.trace_event") as trace:
|
|
manager = ProviderRuntimeManager(
|
|
_settings("nvidia_nim/one"),
|
|
runtime_factory=factory,
|
|
)
|
|
lease = await manager.acquire()
|
|
await manager.replace(
|
|
_settings("nvidia_nim/two"),
|
|
commit=lambda: None,
|
|
reason="test_replace",
|
|
)
|
|
await lease.release()
|
|
await manager.close()
|
|
|
|
events = [call.kwargs for call in trace.call_args_list]
|
|
names = [event["event"] for event in events]
|
|
assert names == [
|
|
"provider_generation.published",
|
|
"provider_generation.published",
|
|
"provider_generation.retired",
|
|
"provider_generation.closed",
|
|
"provider_generation.retired",
|
|
"provider_generation.closed",
|
|
]
|
|
assert events[1]["generation_id"] == 2
|
|
assert events[1]["previous_generation_id"] == 1
|
|
assert events[1]["reason"] == "test_replace"
|
|
assert events[2]["active_leases"] == 1
|
|
assert events[3]["forced"] is False
|