From 2d73e6c8d12da5d907fdc037fa50a2557fc9ce63 Mon Sep 17 00:00:00 2001
From: Simon Willison <swillison@gmail.com>
Date: Sat, 19 Aug 2023 22:02:58 -0700
Subject: [PATCH] Test set_alias(..., alias) too, refs #154

---
 tests/test_aliases.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/tests/test_aliases.py b/tests/test_aliases.py
index 8fc2d8d..94a8ddb 100644
--- a/tests/test_aliases.py
+++ b/tests/test_aliases.py
@@ -5,10 +5,11 @@ import json
 import pytest
 
 
-def test_set_alias():
+@pytest.mark.parametrize("model_id_or_alias", ("gpt-3.5-turbo", "chatgpt"))
+def test_set_alias(model_id_or_alias):
     with pytest.raises(llm.UnknownModelError):
         llm.get_model("this-is-a-new-alias")
-    llm.set_alias("this-is-a-new-alias", "gpt-3.5-turbo")
+    llm.set_alias("this-is-a-new-alias", model_id_or_alias)
     assert llm.get_model("this-is-a-new-alias").model_id == "gpt-3.5-turbo"