From 043a2da1aaac9b48b9c141ebe96603531aebb03d Mon Sep 17 00:00:00 2001 From: hajdul88 <52442977+hajdul88@users.noreply.github.com> Date: Thu, 11 Dec 2025 20:41:01 +0100 Subject: [PATCH] Update test_search_db.py --- cognee/tests/test_search_db.py | 24 ++++++++++++++++++++++++ 1 file changed, 24 insertions(+) diff --git a/cognee/tests/test_search_db.py b/cognee/tests/test_search_db.py index d157a58a9..ba28e9a93 100644 --- a/cognee/tests/test_search_db.py +++ b/cognee/tests/test_search_db.py @@ -1,8 +1,11 @@ import pathlib import os +import warnings 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 from cognee.modules.graph.cognee_graph.CogneeGraphElements import Edge @@ -46,6 +49,27 @@ async def cleanup_litellm_clients(): pass +@pytest_asyncio.fixture(scope="session", autouse=True) +async def cleanup_litellm_clients_session(): + """Session-scoped fixture to cleanup LiteLLM clients at the end of all tests. + + This ensures cleanup happens before LiteLLM's own cleanup tries to run + during event loop closure, preventing RuntimeWarning in Python 3.10. + """ + yield + # Final cleanup of LiteLLM async clients at session end + try: + import litellm + + if hasattr(litellm, "close_litellm_async_clients"): + cleanup_coro = litellm.close_litellm_async_clients() + if cleanup_coro is not None: + await cleanup_coro + except (RuntimeError, Exception): + # Event loop might already be closing, ignore the error + pass + + 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