1 files changed, 26 insertions, 0 deletions
diff --git a/test_llamachat.py b/test_llamachat.py
index 20aea57..bf43c77 100755
--- a/test_llamachat.py
+++ b/test_llamachat.py
@@ -520,6 +520,31 @@ def test_audio_bubble_marker():
     print("ok  audio bubble marker")
 
 
+def test_model_label_markers():
+    """The picker marks vision and audio, and strips both back off."""
+    import os
+
+    os.environ.setdefault("QT_QPA_PLATFORM", "offscreen")
+
+    from llamachat import models
+    from llamachat.ui import _model_label, _strip_marker
+
+    assert _model_label("m", models.ModelInfo()) == "m"
+    assert _model_label("m", models.ModelInfo(vision=True)) == "m 👁"
+    assert _model_label("m", models.ModelInfo(audio=True)) == "m 🎙"
+    assert (
+        _model_label("m", models.ModelInfo(vision=True, audio=True))
+        == "m 👁 🎙"
+    )
+
+    # The strip is the inverse, and leaves an unmarked name alone.
+    assert _strip_marker("m 👁 🎙") == "m"
+    assert _strip_marker("m 🎙") == "m"
+    assert _strip_marker("m 👁") == "m"
+    assert _strip_marker("m") == "m"
+    print("ok  model label markers")
+
+
 def test_unbalanced_backticks():
     """An odd backtick cannot reflow the rest of a reply as code."""
     import os
@@ -4065,6 +4090,7 @@ if __name__ == "__main__":
     test_prompt_column_migration()
     test_markdown_rendering()
     test_audio_bubble_marker()
+    test_model_label_markers()
     test_unbalanced_backticks()
     test_sidebar_toggle()
     test_shortcuts()