Reload pydantic before and after tests

src/tests/llms/test_llms.py

@@ -62,6 +62,12 @@
 OPENAI_TREE_CHAT_ANSWER = "Trees can be various shades of green, but they"
 
 
+def _reload_pydantic():
+    for m in list(sys.modules.keys()):
+        if m.startswith("pydantic"):
+            del sys.modules[m]
+
+
 class TestLLM:
     def test_import(self):
         from concurrent import futures
@@ -2104,13 +2110,12 @@ def setup_class(cls):
         cls.pydantic_version = importlib.metadata.version("pydantic")
         os.system("pip3 install petals==2.2.0")
         os.system("pip3 install 'pydantic>=1.10,<2.0'")
-        for m in list(sys.modules.keys()):
-            if m.startswith("pydantic"):
-                del sys.modules[m]
+        _reload_pydantic()
 
     @classmethod
     def teardown_class(cls):
         os.system(f"pip3 install pydantic=={cls.pydantic_version}")
+        _reload_pydantic()
 
     def test_init(self, create_datadreamer):
         with create_datadreamer():
@@ -2443,10 +2448,12 @@ class TestTogether:
     def setup_class(cls):
         cls.pydantic_version = importlib.metadata.version("pydantic")
         os.system("pip3 install together==0.2.10")
+        _reload_pydantic()
 
     @classmethod
     def teardown_class(cls):
         os.system(f"pip3 install pydantic=={cls.pydantic_version}")
+        _reload_pydantic()
 
     def test_warnings(self, create_datadreamer):
         with create_datadreamer():
@@ -2624,10 +2631,12 @@ class TestMistralAI:
     def setup_class(cls):
         cls.pydantic_version = importlib.metadata.version("pydantic")
         os.system("pip3 install mistralai==0.0.8")
+        _reload_pydantic()
 
     @classmethod
     def teardown_class(cls):
         os.system(f"pip3 install pydantic=={cls.pydantic_version}")
+        _reload_pydantic()
 
     def test_init(self, create_datadreamer):
         with create_datadreamer():