diff --git a/cognee/tests/test_search_db.py b/cognee/tests/test_search_db.py index a20976aaf..51adbf936 100644 --- a/cognee/tests/test_search_db.py +++ b/cognee/tests/test_search_db.py @@ -4,7 +4,6 @@ import warnings import atexit import inspect import pytest -import pytest_asyncio # In Python 3.10, LiteLLM registers an atexit hook (`register_async_client_cleanup`) that # can schedule/instantiate an event loop during interpreter shutdown and still emit: @@ -77,6 +76,7 @@ from collections import Counter logger = get_logger() + 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 @@ -193,8 +193,12 @@ async def setup_test_environment_for_feedback(): @pytest.mark.asyncio -async def test_graph_vector_and_retrieval(): - """Combine graph/vector consistency + retriever checks to avoid repeating expensive setup.""" +async def test_search_db(): + """Run all search-db checks in one test. + + This intentionally keeps everything in a single test to avoid event loop churn between + tests when running against deployed databases. + """ await setup_test_environment() # --- Graph/vector engine consistency --- @@ -273,12 +277,7 @@ async def test_graph_vector_and_retrieval(): f"{name}: node_2 vector_distance {distance} out of [0,1], this shouldn't happen" ) - -@pytest.mark.asyncio -async def test_search_and_graph_side_effects(): - """Combine search result checks + graph side effects to avoid repeating expensive setup.""" - await setup_test_environment() - + # --- Search operations + graph side effects --- completion_gk = await cognee.search( query_type=SearchType.GRAPH_COMPLETION, query_text="Where is germany located, next to which country?", @@ -348,12 +347,14 @@ async def test_search_and_graph_side_effects(): # Assert there are exactly 4 CogneeUserInteraction nodes. assert type_counts.get("CogneeUserInteraction", 0) == 4, ( - f"Expected exactly four CogneeUserInteraction nodes, but found {type_counts.get('CogneeUserInteraction', 0)}" + "Expected exactly four CogneeUserInteraction nodes, " + f"but found {type_counts.get('CogneeUserInteraction', 0)}" ) # Assert there is exactly two CogneeUserFeedback nodes. assert type_counts.get("CogneeUserFeedback", 0) == 2, ( - f"Expected exactly two CogneeUserFeedback nodes, but found {type_counts.get('CogneeUserFeedback', 0)}" + "Expected exactly two CogneeUserFeedback nodes, " + f"but found {type_counts.get('CogneeUserFeedback', 0)}" ) # Assert there is exactly two NodeSet. @@ -363,17 +364,20 @@ async def test_search_and_graph_side_effects(): # Assert that there are at least 10 'used_graph_element_to_answer' edges. assert edge_type_counts.get("used_graph_element_to_answer", 0) >= 10, ( - f"Expected at least ten 'used_graph_element_to_answer' edges, but found {edge_type_counts.get('used_graph_element_to_answer', 0)}" + "Expected at least ten 'used_graph_element_to_answer' edges, but found " + f"{edge_type_counts.get('used_graph_element_to_answer', 0)}" ) # Assert that there are exactly 2 'gives_feedback_to' edges. assert edge_type_counts.get("gives_feedback_to", 0) == 2, ( - f"Expected exactly two 'gives_feedback_to' edges, but found {edge_type_counts.get('gives_feedback_to', 0)}" + "Expected exactly two 'gives_feedback_to' edges, but found " + f"{edge_type_counts.get('gives_feedback_to', 0)}" ) # Assert that there are at least 6 'belongs_to_set' edges. assert edge_type_counts.get("belongs_to_set", 0) >= 6, ( - f"Expected at least six 'belongs_to_set' edges, but found {edge_type_counts.get('belongs_to_set', 0)}" + "Expected at least six 'belongs_to_set' edges, but found " + f"{edge_type_counts.get('belongs_to_set', 0)}" ) # Node field validation on the same graph produced above @@ -385,7 +389,8 @@ async def test_search_and_graph_side_effects(): for node_id, data in nodes: if data.get("type") == "CogneeUserInteraction": assert required_fields_user_interaction.issubset(data.keys()), ( - f"Node {node_id} is missing fields: {required_fields_user_interaction - set(data.keys())}" + f"Node {node_id} is missing fields: " + f"{required_fields_user_interaction - set(data.keys())}" ) for field in required_fields_user_interaction: @@ -405,10 +410,7 @@ async def test_search_and_graph_side_effects(): f"Node {node_id} has invalid value for '{field}': {value!r}" ) - -@pytest.mark.asyncio -async def test_feedback_weight_calculation(): - """Test that feedback weight is correctly calculated after multiple positive feedbacks.""" + # --- Feedback weight calculation (run in fresh environment) --- await setup_test_environment_for_feedback() await cognee.search( @@ -438,3 +440,26 @@ async def test_feedback_weight_calculation(): assert properties["feedback_weight"] >= 6, ( "Feedback weight calculation is not correct, it should be more then 6." ) + + # Best-effort cleanup (helps avoid CI shutdown warnings/crashes). + try: + vector_engine = get_vector_engine() + if hasattr(vector_engine, "engine") and hasattr(vector_engine.engine, "dispose"): + await vector_engine.engine.dispose(close=True) + except Exception: + pass + + 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 Exception: + pass + + try: + _remove_litellm_atexit_handlers() + except Exception: + pass