gmakstutis/cognee
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Update test_search_db.py

Browse source
This commit is contained in:
hajdul88 2025-12-11 20:06:30 +01:00
parent b2d2fd971d
commit 41a3d75ed8
1 changed files with 18 additions and 14 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

18
cognee/tests/test_search_db.py
View file

@ -28,19 +28,21 @@ async def setup_test_environment():
# This test runs for multiple db settings, to run this locally set the corresponding db envs # This test runs for multiple db settings, to run this locally set the corresponding db envs
# Dispose of existing engines and clear caches to ensure fresh instances for each test # Dispose of existing engines and clear caches to ensure fresh instances for each test
# This prevents event loop issues when using deployed databases (Neo4j, PostgreSQL)
try: try:
from cognee.infrastructure.databases.vector import get_vector_engine from cognee.infrastructure.databases.vector import get_vector_engine
vector_engine = get_vector_engine() vector_engine = get_vector_engine()
# Dispose SQLAlchemy engine connection pool if it exists # Dispose SQLAlchemy engine connection pool if it exists
if hasattr(vector_engine, 'engine') and hasattr(vector_engine.engine, 'dispose'): if hasattr(vector_engine, "engine") and hasattr(vector_engine.engine, "dispose"):
await vector_engine.engine.dispose(close=True) await vector_engine.engine.dispose(close=True)
except Exception: except Exception:
pass # Engine might not exist yet pass
from cognee.infrastructure.databases.graph.get_graph_engine import create_graph_engine from cognee.infrastructure.databases.graph.get_graph_engine import create_graph_engine
from cognee.infrastructure.databases.vector.create_vector_engine import create_vector_engine from cognee.infrastructure.databases.vector.create_vector_engine import create_vector_engine
from cognee.infrastructure.databases.relational.create_relational_engine import create_relational_engine from cognee.infrastructure.databases.relational.create_relational_engine import (
create_relational_engine,
)
create_graph_engine.cache_clear() create_graph_engine.cache_clear()
create_vector_engine.cache_clear() create_vector_engine.cache_clear()
@ -95,16 +97,19 @@ async def setup_test_environment_for_feedback():
# This prevents event loop issues when using deployed databases (Neo4j, PostgreSQL) # This prevents event loop issues when using deployed databases (Neo4j, PostgreSQL)
try: try:
from cognee.infrastructure.databases.vector import get_vector_engine from cognee.infrastructure.databases.vector import get_vector_engine
vector_engine = get_vector_engine() vector_engine = get_vector_engine()
# Dispose SQLAlchemy engine connection pool if it exists # Dispose SQLAlchemy engine connection pool if it exists
if hasattr(vector_engine, 'engine') and hasattr(vector_engine.engine, 'dispose'): if hasattr(vector_engine, "engine") and hasattr(vector_engine.engine, "dispose"):
await vector_engine.engine.dispose(close=True) await vector_engine.engine.dispose(close=True)
except Exception: except Exception:
pass # Engine might not exist yet pass # Engine might not exist yet
from cognee.infrastructure.databases.graph.get_graph_engine import create_graph_engine from cognee.infrastructure.databases.graph.get_graph_engine import create_graph_engine
from cognee.infrastructure.databases.vector.create_vector_engine import create_vector_engine from cognee.infrastructure.databases.vector.create_vector_engine import create_vector_engine
from cognee.infrastructure.databases.relational.create_relational_engine import create_relational_engine from cognee.infrastructure.databases.relational.create_relational_engine import (
create_relational_engine,
)
create_graph_engine.cache_clear() create_graph_engine.cache_clear()
create_vector_engine.cache_clear() create_vector_engine.cache_clear()
@ -151,7 +156,6 @@ async def test_graph_vector_engine_consistency():
) )
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_retriever_contexts(): async def test_retriever_contexts():
"""Test that all retrievers return valid contexts with expected content.""" """Test that all retrievers return valid contexts with expected content."""