| 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.application.ports import RequestRuntimePort |
| 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 |
|
|
|
|
| class RecordingModelCatalogPublisher: |
| def __init__(self) -> None: |
| self.events: list[str] = [] |
| self.snapshots: list[tuple[str, tuple[str, ...]]] = [] |
|
|
| def ensure_exists(self, runtime: RequestRuntimePort) -> None: |
| self._record("ensure_exists", runtime) |
|
|
| def publish(self, runtime: RequestRuntimePort) -> None: |
| self._record("publish", runtime) |
|
|
| def _record(self, event: str, runtime: RequestRuntimePort) -> None: |
| self.events.append(event) |
| self.snapshots.append( |
| ( |
| runtime.current_settings().model, |
| tuple(info.model_id for info in runtime.cached_prefixed_model_infos()), |
| ) |
| ) |
|
|
|
|
| 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_connected_provider_refresh_and_eviction_are_targeted() -> None: |
| factory = RuntimeFactory() |
| connected: set[str] = set() |
| manager = ProviderRuntimeManager( |
| _settings("nvidia_nim/one"), |
| runtime_factory=factory, |
| connected_provider_ids=lambda: tuple(connected), |
| ) |
| factory.runtimes[0].provider.list_model_infos = AsyncMock( |
| return_value=frozenset( |
| {ProviderModelInfo(model_id="gpt-visible", supports_thinking=True)} |
| ) |
| ) |
|
|
| connected.add("openai") |
| result = await manager.connected_provider_changed("openai", connected=True) |
|
|
| assert result.refreshed_provider_ids == ("openai",) |
| assert manager.cached_model_ids()["openai"] == frozenset({"gpt-visible"}) |
|
|
| connected.clear() |
| assert "openai" not in manager.cached_model_ids() |
|
|
| await manager.connected_provider_changed("openai", connected=False) |
|
|
| assert "openai" not in manager.cached_model_ids() |
| await manager.close() |
|
|
|
|
| @pytest.mark.asyncio |
| async def test_catalog_publication_tracks_warm_refresh_and_direct_cache() -> None: |
| factory = RuntimeFactory() |
| publisher = RecordingModelCatalogPublisher() |
| manager = ProviderRuntimeManager( |
| _settings("nvidia_nim/one"), |
| runtime_factory=factory, |
| model_catalog_publisher=publisher, |
| ) |
| factory.runtimes[0].provider.list_model_infos = AsyncMock( |
| return_value=frozenset({ProviderModelInfo("warm-model")}) |
| ) |
|
|
| await manager.warm_referenced_model_cache() |
| manager.start_model_list_refresh() |
| refresh_task = manager._refresh_task |
| assert refresh_task is not None |
| await refresh_task |
| manager.cache_model_infos( |
| "nvidia_nim", |
| {ProviderModelInfo("tested-model")}, |
| ) |
|
|
| assert publisher.events == ["ensure_exists", "publish", "publish"] |
| assert publisher.snapshots == [ |
| ("nvidia_nim/one", ("nvidia_nim/warm-model",)), |
| ("nvidia_nim/one", ("nvidia_nim/warm-model",)), |
| ("nvidia_nim/one", ("nvidia_nim/tested-model",)), |
| ] |
|
|
| await manager.close() |
| assert publisher.events == ["ensure_exists", "publish", "publish"] |
|
|
|
|
| @pytest.mark.asyncio |
| async def test_failed_startup_discovery_still_ensures_a_fresh_catalog() -> None: |
| factory = RuntimeFactory() |
| publisher = RecordingModelCatalogPublisher() |
| manager = ProviderRuntimeManager( |
| _settings("nvidia_nim/configured"), |
| runtime_factory=factory, |
| model_catalog_publisher=publisher, |
| ) |
| factory.runtimes[0].provider.list_model_infos = AsyncMock( |
| side_effect=RuntimeError("upstream unavailable") |
| ) |
|
|
| result = await manager.warm_referenced_model_cache() |
|
|
| assert result.failed_provider_ids == ("nvidia_nim",) |
| assert publisher.events == ["ensure_exists"] |
| assert publisher.snapshots == [("nvidia_nim/configured", ())] |
| await manager.close() |
|
|
|
|
| @pytest.mark.asyncio |
| async def test_catalog_publication_tracks_connected_account_changes() -> None: |
| factory = RuntimeFactory() |
| publisher = RecordingModelCatalogPublisher() |
| connected: set[str] = set() |
| manager = ProviderRuntimeManager( |
| _settings("nvidia_nim/one"), |
| runtime_factory=factory, |
| connected_provider_ids=lambda: tuple(connected), |
| model_catalog_publisher=publisher, |
| ) |
| factory.runtimes[0].provider.list_model_infos = AsyncMock( |
| return_value=frozenset({ProviderModelInfo("gpt-connected")}) |
| ) |
|
|
| connected.add("openai") |
| await manager.connected_provider_changed("openai", connected=True) |
| connected.clear() |
| await manager.connected_provider_changed("openai", connected=False) |
|
|
| assert publisher.events == ["publish", "publish"] |
| assert publisher.snapshots == [ |
| ("nvidia_nim/one", ("openai/gpt-connected",)), |
| ("nvidia_nim/one", ()), |
| ] |
| await manager.close() |
|
|
|
|
| @pytest.mark.asyncio |
| async def test_catalog_publication_tracks_replacement_and_its_refresh() -> None: |
| factory = RuntimeFactory() |
| publisher = RecordingModelCatalogPublisher() |
| manager = ProviderRuntimeManager( |
| _settings("nvidia_nim/one"), |
| runtime_factory=factory, |
| model_catalog_publisher=publisher, |
| ) |
|
|
| await manager.replace( |
| _settings("nvidia_nim/two"), |
| commit=lambda: None, |
| ) |
| refresh_task = manager._refresh_task |
| assert refresh_task is not None |
| await refresh_task |
|
|
| assert publisher.events == ["publish", "publish"] |
| assert publisher.snapshots == [ |
| ("nvidia_nim/two", ()), |
| ("nvidia_nim/two", ()), |
| ] |
| await manager.close() |
|
|
|
|
| @pytest.mark.asyncio |
| async def test_catalog_publication_failure_is_warning_only_and_secret_safe() -> None: |
| factory = RuntimeFactory() |
| secret = "private-catalog-write-detail" |
| publisher = MagicMock() |
| publisher.ensure_exists.side_effect = PermissionError(secret) |
| publisher.publish.side_effect = PermissionError(secret) |
| manager = ProviderRuntimeManager( |
| _settings("nvidia_nim/one"), |
| runtime_factory=factory, |
| model_catalog_publisher=publisher, |
| ) |
|
|
| with patch("free_claude_code.runtime.provider_manager.logger.warning") as warning: |
| await manager.warm_referenced_model_cache() |
| manager.cache_model_infos( |
| "nvidia_nim", |
| {ProviderModelInfo("tested-model")}, |
| ) |
|
|
| assert warning.call_count == 2 |
| log_blob = " ".join( |
| str(value) |
| for call in warning.call_args_list |
| for value in (*call.args, *call.kwargs.values()) |
| ) |
| assert "PermissionError" in log_blob |
| assert secret not in log_blob |
| await manager.close() |
|
|
|
|
| @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_hot_replacement_owns_admission_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._admission is not old_provider._admission |
| 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("lmstudio", {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() == {"lmstudio": 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("lmstudio", {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() == {"lmstudio": 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("nvidia_nim/one"), |
| runtime_factory=factory, |
| ) |
| manager.cache_model_infos( |
| "lmstudio", |
| {ProviderModelInfo("persisted", supports_thinking=True)}, |
| ) |
|
|
| await manager.replace( |
| _settings("nvidia_nim/two"), |
| commit=lambda: None, |
| ) |
|
|
| assert manager.cached_model_ids() == {"lmstudio": frozenset({"persisted"})} |
| assert manager.cached_model_supports_thinking("lmstudio", "persisted") is True |
| await manager.close() |
|
|
|
|
| @pytest.mark.asyncio |
| async def test_replacement_prunes_and_rejects_removed_remote_provider_cache() -> None: |
| factory = RuntimeFactory() |
| first_settings = _settings("open_router/one").model_copy( |
| update={"open_router_api_key": "open-router-key"} |
| ) |
| manager = ProviderRuntimeManager(first_settings, runtime_factory=factory) |
| manager.cache_model_infos("open_router", {ProviderModelInfo("old-model")}) |
|
|
| await manager.replace( |
| _settings("nvidia_nim/two"), |
| commit=lambda: None, |
| ) |
| manager.cache_model_infos("open_router", {ProviderModelInfo("late-old-model")}) |
|
|
| assert "open_router" not in manager.cached_model_ids() |
| 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 |
|
|