test_image_gen_picker.py
1 """Tests for plugin image_gen providers injecting themselves into the picker. 2 3 Covers `_plugin_image_gen_providers`, `_visible_providers`, and 4 `_toolset_needs_configuration_prompt` handling of plugin providers. 5 """ 6 7 from __future__ import annotations 8 9 from types import SimpleNamespace 10 11 import pytest 12 13 from agent import image_gen_registry 14 from agent.image_gen_provider import ImageGenProvider 15 16 17 class _FakeProvider(ImageGenProvider): 18 def __init__(self, name: str, available: bool = True, schema=None, models=None): 19 self._name = name 20 self._available = available 21 self._schema = schema or { 22 "name": name.title(), 23 "badge": "test", 24 "tag": f"{name} test tag", 25 "env_vars": [{"key": f"{name.upper()}_API_KEY", "prompt": f"{name} key"}], 26 } 27 self._models = models or [ 28 {"id": f"{name}-model-v1", "display": f"{name} v1", 29 "speed": "~5s", "strengths": "test", "price": "$"}, 30 ] 31 32 @property 33 def name(self) -> str: 34 return self._name 35 36 def is_available(self) -> bool: 37 return self._available 38 39 def list_models(self): 40 return list(self._models) 41 42 def default_model(self): 43 return self._models[0]["id"] if self._models else None 44 45 def get_setup_schema(self): 46 return dict(self._schema) 47 48 def generate(self, prompt, aspect_ratio="landscape", **kw): 49 return {"success": True, "image": f"{self._name}://{prompt}"} 50 51 52 @pytest.fixture(autouse=True) 53 def _reset_registry(): 54 image_gen_registry._reset_for_tests() 55 yield 56 image_gen_registry._reset_for_tests() 57 58 59 class TestPluginPickerInjection: 60 def test_plugin_providers_returns_registered(self, monkeypatch): 61 from hermes_cli import tools_config 62 63 image_gen_registry.register_provider(_FakeProvider("myimg")) 64 65 rows = tools_config._plugin_image_gen_providers() 66 names = [r["name"] for r in rows] 67 plugin_names = [r.get("image_gen_plugin_name") for r in rows] 68 69 assert "Myimg" in names 70 assert "myimg" in plugin_names 71 72 def test_fal_skipped_to_avoid_duplicate(self, monkeypatch): 73 from hermes_cli import tools_config 74 75 # Simulate a FAL plugin being registered — the picker already has 76 # hardcoded FAL rows in TOOL_CATEGORIES, so plugin-FAL must be 77 # skipped to avoid showing FAL twice. 78 image_gen_registry.register_provider(_FakeProvider("fal")) 79 image_gen_registry.register_provider(_FakeProvider("openai")) 80 81 rows = tools_config._plugin_image_gen_providers() 82 names = [r.get("image_gen_plugin_name") for r in rows] 83 assert "fal" not in names 84 assert "openai" in names 85 86 def test_visible_providers_includes_plugins_for_image_gen(self, monkeypatch): 87 from hermes_cli import tools_config 88 89 image_gen_registry.register_provider(_FakeProvider("someimg")) 90 91 cat = tools_config.TOOL_CATEGORIES["image_gen"] 92 visible = tools_config._visible_providers(cat, {}) 93 plugin_names = [p.get("image_gen_plugin_name") for p in visible if p.get("image_gen_plugin_name")] 94 assert "someimg" in plugin_names 95 96 def test_visible_providers_does_not_inject_into_other_categories(self, monkeypatch): 97 from hermes_cli import tools_config 98 99 image_gen_registry.register_provider(_FakeProvider("someimg")) 100 101 # Browser category must NOT see image_gen plugins. 102 browser = tools_config.TOOL_CATEGORIES["browser"] 103 visible = tools_config._visible_providers(browser, {}) 104 assert all(p.get("image_gen_plugin_name") is None for p in visible) 105 106 107 class TestPluginCatalog: 108 def test_plugin_catalog_returns_models(self): 109 from hermes_cli import tools_config 110 111 image_gen_registry.register_provider(_FakeProvider("catimg")) 112 113 catalog, default = tools_config._plugin_image_gen_catalog("catimg") 114 assert "catimg-model-v1" in catalog 115 assert default == "catimg-model-v1" 116 117 def test_plugin_catalog_empty_for_unknown(self): 118 from hermes_cli import tools_config 119 120 catalog, default = tools_config._plugin_image_gen_catalog("does-not-exist") 121 assert catalog == {} 122 assert default is None 123 124 125 class TestConfigPrompt: 126 def test_image_gen_satisfied_by_plugin_provider(self, monkeypatch, tmp_path): 127 """When a plugin provider reports is_available(), the picker should 128 not force a setup prompt on the user.""" 129 from hermes_cli import tools_config 130 131 monkeypatch.setenv("HERMES_HOME", str(tmp_path)) 132 monkeypatch.delenv("FAL_KEY", raising=False) 133 134 image_gen_registry.register_provider(_FakeProvider("avail-img", available=True)) 135 136 assert tools_config._toolset_needs_configuration_prompt("image_gen", {}) is False 137 138 def test_image_gen_still_prompts_when_nothing_available(self, monkeypatch, tmp_path): 139 from hermes_cli import tools_config 140 141 monkeypatch.setenv("HERMES_HOME", str(tmp_path)) 142 monkeypatch.delenv("FAL_KEY", raising=False) 143 144 image_gen_registry.register_provider(_FakeProvider("unavail-img", available=False)) 145 146 assert tools_config._toolset_needs_configuration_prompt("image_gen", {}) is True 147 148 149 class TestConfigWriting: 150 def test_picking_plugin_provider_writes_provider_and_model(self, monkeypatch, tmp_path): 151 """When a user picks a plugin-backed image_gen provider with no 152 env vars needed, ``_configure_provider`` should write both 153 ``image_gen.provider`` and ``image_gen.model``.""" 154 from hermes_cli import tools_config 155 156 monkeypatch.setenv("HERMES_HOME", str(tmp_path)) 157 image_gen_registry.register_provider(_FakeProvider("noenv", schema={ 158 "name": "NoEnv", 159 "badge": "free", 160 "tag": "", 161 "env_vars": [], 162 })) 163 164 # Stub out the interactive model picker — no TTY in tests. 165 monkeypatch.setattr(tools_config, "_prompt_choice", lambda *a, **kw: 0) 166 167 config: dict = {} 168 provider_row = { 169 "name": "NoEnv", 170 "env_vars": [], 171 "image_gen_plugin_name": "noenv", 172 } 173 tools_config._configure_provider(provider_row, config) 174 175 assert config["image_gen"]["provider"] == "noenv" 176 assert config["image_gen"]["model"] == "noenv-model-v1" 177 178 def test_reconfiguring_plugin_provider_writes_provider_and_model(self, monkeypatch, tmp_path): 179 """The reconfigure path should switch image_gen away from managed FAL 180 and onto the selected plugin provider.""" 181 from hermes_cli import tools_config 182 183 monkeypatch.setenv("HERMES_HOME", str(tmp_path)) 184 image_gen_registry.register_provider(_FakeProvider("testopenai")) 185 monkeypatch.setattr(tools_config, "_prompt_choice", lambda *a, **kw: 0) 186 monkeypatch.setattr(tools_config, "_prompt", lambda *a, **kw: "") 187 monkeypatch.setattr( 188 tools_config, 189 "get_env_value", 190 lambda key: "sk-test" if key == "OPENAI_API_KEY" else "", 191 ) 192 193 config = {"image_gen": {"use_gateway": True}} 194 provider_row = { 195 "name": "OpenAI", 196 "env_vars": [{"key": "OPENAI_API_KEY", "prompt": "OpenAI API key"}], 197 "image_gen_plugin_name": "testopenai", 198 } 199 200 tools_config._reconfigure_provider(provider_row, config) 201 202 assert config["image_gen"]["provider"] == "testopenai" 203 assert config["image_gen"]["model"] == "testopenai-model-v1" 204 assert config["image_gen"]["use_gateway"] is False 205 206 def test_plugin_provider_active_overrides_managed_nous_active_label(self, monkeypatch): 207 from hermes_cli import tools_config 208 209 monkeypatch.setattr( 210 tools_config, 211 "get_nous_subscription_features", 212 lambda config: SimpleNamespace( 213 features={"image_gen": SimpleNamespace(managed_by_nous=True)} 214 ), 215 ) 216 217 config = {"image_gen": {"provider": "openai", "use_gateway": False}} 218 nous_row = { 219 "name": "Nous Subscription", 220 "managed_nous_feature": "image_gen", 221 } 222 openai_row = { 223 "name": "OpenAI", 224 "image_gen_plugin_name": "openai", 225 } 226 227 assert tools_config._is_provider_active(openai_row, config) is True 228 assert tools_config._is_provider_active(nous_row, config) is False 229 230 def test_reconfiguring_fal_clears_plugin_provider(self, monkeypatch): 231 from hermes_cli import tools_config 232 233 monkeypatch.setattr(tools_config, "_prompt_choice", lambda *a, **kw: 0) 234 monkeypatch.setattr(tools_config, "_prompt", lambda *a, **kw: "") 235 monkeypatch.setattr( 236 tools_config, 237 "get_env_value", 238 lambda key: "fal-key" if key == "FAL_KEY" else "", 239 ) 240 241 config = {"image_gen": {"provider": "openai", "use_gateway": False}} 242 provider_row = { 243 "name": "FAL.ai", 244 "env_vars": [{"key": "FAL_KEY", "prompt": "FAL API key"}], 245 "imagegen_backend": "fal", 246 } 247 248 tools_config._reconfigure_provider(provider_row, config) 249 250 assert config["image_gen"]["provider"] == "fal" 251 assert config["image_gen"]["use_gateway"] is False