From 2cd2b20adcd46f8e68174bd21999ddb1e62ab889 Mon Sep 17 00:00:00 2001 From: Ajay Patel Date: Mon, 25 Dec 2023 22:57:05 -0500 Subject: [PATCH] Reload pydantic before and after tests --- src/tests/llms/test_llms.py | 15 ++++++++++++--- 1 file changed, 12 insertions(+), 3 deletions(-) diff --git a/src/tests/llms/test_llms.py b/src/tests/llms/test_llms.py index c255fc9e..1fbb0461 100644 --- a/src/tests/llms/test_llms.py +++ b/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():