From 7ff849224fb85bfba27d19f4cdf427281690d40c Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Erik=20Bj=C3=A4reholt?= <erik@bjareho.lt>
Date: Sat, 21 Dec 2024 17:41:10 +0100
Subject: [PATCH] test: reset default model in openai tests to fix flaky test

---
 gptme/llm/models.py      |  4 ++--
 tests/test_llm_openai.py | 11 ++++++++++-
 2 files changed, 12 insertions(+), 3 deletions(-)

diff --git a/gptme/llm/models.py b/gptme/llm/models.py
index 1a2bd76b..d3e42f33 100644
--- 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
index eb3f9b1d..d2addfd4 100644
--- 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),