diff --git a/cognee/tests/test_search_db.py b/cognee/tests/test_search_db.py index a93424176..a94d029a3 100644 --- a/cognee/tests/test_search_db.py +++ b/cognee/tests/test_search_db.py @@ -1,6 +1,7 @@ import pathlib import os import pytest +import pytest_asyncio import cognee from cognee.infrastructure.databases.graph import get_graph_engine from cognee.infrastructure.databases.vector import get_vector_engine @@ -23,6 +24,19 @@ from collections import Counter logger = get_logger() +@pytest_asyncio.fixture(autouse=True) +async def cleanup_litellm_clients(): + """Fixture to properly cleanup LiteLLM async clients after each test.""" + yield + # Cleanup LiteLLM async clients to prevent RuntimeWarning about unawaited coroutine + try: + import litellm + if hasattr(litellm, "close_litellm_async_clients"): + await litellm.close_litellm_async_clients() + except Exception: + pass # LiteLLM might not be available or already cleaned up + + async def setup_test_environment(): """Helper function to set up test environment with data, cognify, and triplet embeddings.""" # This test runs for multiple db settings, to run this locally set the corresponding db envs