added image generation

3 files changed, 43 insertions, 0 deletions

diff --git a/tests/context_test.py b/tests/context_test.py
index a1d9624..7cb4810 100644
--- a/tests/context_test.py
+++ b/tests/context_test.py
@@ -24,6 +24,23 @@ default_config = {
   },
 }
 
+default_image_config = {
+  "prompt": "",
+  "options": {
+    "model": "dall-e-3",
+    "endpoint_url": "https://api.openai.com/v1/images/generations",
+    "quality": "standard",
+    "size": "1024x1024",
+    "style": "vivid",
+    "request_timeout": "20",
+    "enable_auth": "1",
+    "token_file_path": "",
+  },
+  "ui": {
+    "paste_mode": "1",
+  },
+}
+
 def test_default_config():
     actual_context = make_ai_context({
         'config_default': default_config,
@@ -130,6 +147,25 @@ def test_chat_only_role():
     actual_config = context['config']
     assert 'preset_tab' == actual_config['options']['open_chat_command']
 
+def test_image_role():
+    actual_context = make_ai_context({
+        'config_default': default_image_config,
+        'config_extension': {},
+        'user_instruction': '/hd-image picture of the moon',
+        'user_selection': '',
+        'command_type': 'image',
+    })
+    expected_options = {
+        **default_image_config['options'],
+        'token_file_path': '/custom/path/ai.token',
+        'quality': 'hd',
+    }
+    expected_context = {
+        'config': { **default_image_config, 'options': expected_options },
+        'prompt': 'picture of the moon',
+    }
+    assert expected_context == actual_context
+
 def test_default_roles():
     base = {
         'config_default': default_config,
diff --git a/tests/resources/roles.ini b/tests/resources/roles.ini
index 00903f7..560e06d 100644
--- a/tests/resources/roles.ini
+++ b/tests/resources/roles.ini
@@ -21,6 +21,9 @@ options.endpoint_url = https://localhost/edit
 [chat-only-role.chat]
 options.open_chat_command = preset_tab
 
+[hd-image.image]
+options.quality = hd
+
 [deprecated-test-role-simple]
 prompt = simple role prompt
 [deprecated-test-role-simple.options]
diff --git a/tests/roles_test.py b/tests/roles_test.py
index da9b9a4..a0cc26a 100644
--- a/tests/roles_test.py
+++ b/tests/roles_test.py
@@ -22,3 +22,7 @@ def test_role_chat_only():
         'below',
         'tab',
     }
+
+def test_explicit_image_roles():
+    role_names = load_ai_role_names('image')
+    assert set(role_names) == { 'hd-image' }