104 lines
3.5 KiB
Python
104 lines
3.5 KiB
Python
"""Image generation runtime reload tests."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
from types import SimpleNamespace
|
|
|
|
import pytest
|
|
|
|
from nanobot.agent import context as agent_context
|
|
from nanobot.agent.tools.image_generation import (
|
|
ImageGenerationTool,
|
|
reload_image_generation_tool,
|
|
request_image_generation_reload,
|
|
)
|
|
from nanobot.agent.tools.registry import ToolRegistry
|
|
from nanobot.bus.queue import MessageBus
|
|
from nanobot.config.loader import load_config, save_config
|
|
from nanobot.config.schema import ToolsConfig
|
|
|
|
|
|
def _runtime_state(tmp_path):
|
|
return SimpleNamespace(
|
|
workspace=tmp_path,
|
|
tools_config=ToolsConfig(),
|
|
_image_generation_provider_configs={},
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_image_generation_reload_replaces_and_removes_live_tool(
|
|
tmp_path,
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
config_path = tmp_path / "config.json"
|
|
monkeypatch.setattr("nanobot.config.loader._current_config_path", config_path)
|
|
config = load_config()
|
|
config.providers.openrouter.api_key = "first-key"
|
|
config.tools.image_generation.enabled = True
|
|
config.tools.image_generation.provider = "openrouter"
|
|
config.tools.image_generation.model = "openai/first-image-model"
|
|
save_config(config)
|
|
|
|
state = _runtime_state(tmp_path)
|
|
registry = ToolRegistry()
|
|
result = await reload_image_generation_tool(state, registry)
|
|
|
|
first_tool = registry.get("generate_image")
|
|
assert result["requires_restart"] is False
|
|
assert isinstance(first_tool, ImageGenerationTool)
|
|
assert first_tool.config.model == "openai/first-image-model"
|
|
assert first_tool.provider_configs["openrouter"].api_key == "first-key"
|
|
|
|
config = load_config()
|
|
config.providers.openrouter.api_key = "second-key"
|
|
config.tools.image_generation.model = "openai/second-image-model"
|
|
save_config(config)
|
|
|
|
result = await reload_image_generation_tool(state, registry)
|
|
second_tool = registry.get("generate_image")
|
|
assert result["requires_restart"] is False
|
|
assert isinstance(second_tool, ImageGenerationTool)
|
|
assert second_tool is not first_tool
|
|
assert second_tool.config.model == "openai/second-image-model"
|
|
assert second_tool.provider_configs["openrouter"].api_key == "second-key"
|
|
assert state.tools_config.image_generation.model == "openai/second-image-model"
|
|
|
|
config = load_config()
|
|
config.tools.image_generation.enabled = False
|
|
save_config(config)
|
|
|
|
result = await reload_image_generation_tool(state, registry)
|
|
assert result["requires_restart"] is False
|
|
assert not registry.has("generate_image")
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_image_generation_reload_reaches_agent_runtime_control(
|
|
tmp_path,
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
config_path = tmp_path / "config.json"
|
|
monkeypatch.setattr("nanobot.config.loader._current_config_path", config_path)
|
|
config = load_config()
|
|
config.providers.openrouter.api_key = "image-key"
|
|
config.tools.image_generation.enabled = True
|
|
save_config(config)
|
|
|
|
bus = MessageBus()
|
|
state = _runtime_state(tmp_path)
|
|
registry = ToolRegistry()
|
|
|
|
async def consume_control() -> None:
|
|
message = await bus.consume_inbound()
|
|
assert await agent_context.handle_runtime_control(state, message, registry) is True
|
|
|
|
consumer = asyncio.create_task(consume_control())
|
|
result = await request_image_generation_reload(bus, timeout=2.0)
|
|
await consumer
|
|
|
|
assert result["ok"] is True
|
|
assert result["requires_restart"] is False
|
|
assert registry.has("generate_image")
|