command-type only roles

1 files changed, 13 insertions, 3 deletions

diff --git a/tests/roles_test.py b/tests/roles_test.py
index 3230329..ac5525d 100644
--- a/tests/roles_test.py
+++ b/tests/roles_test.py
@@ -1,10 +1,20 @@
 from roles import load_ai_role_names
 
 def test_role_completion():
-    role_names = load_ai_role_names()
-    assert role_names == [
+    role_names = load_ai_role_names('complete')
+    assert set(role_names) == {
         'test-role-simple',
         'test-role',
         'deprecated-test-role-simple',
         'deprecated-test-role',
-    ]
+    }
+
+def test_role_chat_only():
+    role_names = load_ai_role_names('chat')
+    assert set(role_names) == {
+        'test-role-simple',
+        'test-role',
+        'chat-only-role',
+        'deprecated-test-role-simple',
+        'deprecated-test-role',
+    }