test cleanup

Author: alexander-alderman-webb

tests/integrations/litellm/test_litellm.py

@@ -35,7 +35,9 @@ async def __call__(self, *args, **kwargs):
 
 from concurrent.futures import ThreadPoolExecutor
 
-import litellm.litellm_core_utils.streaming_handler as streaming_handler
+import litellm.utils as litellm_utils
+from litellm.litellm_core_utils import streaming_handler
+from litellm.litellm_core_utils import thread_pool_executor
 
 
 LITELLM_VERSION = package_version("litellm")
@@ -44,7 +46,9 @@ async def __call__(self, *args, **kwargs):
 @pytest.fixture()
 def reset_litellm_executor():
     yield
-    streaming_handler.executor = ThreadPoolExecutor(max_workers=100)
+    thread_pool_executor.executor = ThreadPoolExecutor(max_workers=100)
+    litellm_utils.executor = thread_pool_executor.executor
+    streaming_handler.executor = thread_pool_executor.executor
 
 
 @pytest.fixture