test: reset default model in openai tests to fix flaky test

ErikBjare · Dec 21, 2024 · 7ff8492 · 7ff8492
1 parent 5076376
commit 7ff8492
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 3 deletions.
diff --git a/gptme/llm/models.py b/gptme/llm/models.py
@@ -160,8 +160,8 @@ def get_default_model() -> ModelMeta | None:
     return DEFAULT_MODEL
 
 
-def set_default_model(model: str) -> None:
-    modelmeta = get_model(model)
+def set_default_model(model: str | ModelMeta) -> None:
+    modelmeta = model if isinstance(model, ModelMeta) else get_model(model)
     assert modelmeta
     global DEFAULT_MODEL
     DEFAULT_MODEL = modelmeta

diff --git a/tests/test_llm_openai.py b/tests/test_llm_openai.py
@@ -1,9 +1,18 @@
+import pytest
 from gptme.llm.llm_openai import _prepare_messages_for_api
-from gptme.llm.models import set_default_model
+from gptme.llm.models import get_default_model, set_default_model
 from gptme.message import Message
 from gptme.tools import get_tool, init_tools
 
 
+@pytest.fixture(autouse=True)
+def reset_default_model():
+    default_model = get_default_model()
+    assert default_model
+    yield
+    set_default_model(default_model)
+
+
 def test_message_conversion():
     messages = [
         Message(role="system", content="Initial Message", pinned=True, hide=True),