converts search_db test to pytest test

This commit is contained in:
hajdul88 2025-12-11 17:09:48 +01:00
parent 7a82bd7f7f
commit 3defb9ad44
2 changed files with 8 additions and 12 deletions

View file

@ -54,7 +54,7 @@ jobs:
GRAPH_DATABASE_PROVIDER: 'kuzu'
VECTOR_DB_PROVIDER: 'lancedb'
DB_PROVIDER: 'sqlite'
run: uv run python ./cognee/tests/test_search_db.py
run: uv run pytest cognee/tests/test_search_db.py -v
run-neo4j-lance-sqlite-search-tests:
name: Search test for Neo4j/LanceDB/Sqlite (Python ${{ matrix.python-version }})
@ -100,7 +100,7 @@ jobs:
GRAPH_DATABASE_URL: ${{ steps.neo4j.outputs.neo4j-url }}
GRAPH_DATABASE_USERNAME: ${{ steps.neo4j.outputs.neo4j-username }}
GRAPH_DATABASE_PASSWORD: ${{ steps.neo4j.outputs.neo4j-password }}
run: uv run python ./cognee/tests/test_search_db.py
run: uv run pytest cognee/tests/test_search_db.py -v
run-kuzu-pgvector-postgres-search-tests:
name: Search test for Kuzu/PGVector/Postgres (Python ${{ matrix.python-version }})
@ -159,7 +159,7 @@ jobs:
DB_PORT: 5432
DB_USERNAME: cognee
DB_PASSWORD: cognee
run: uv run python ./cognee/tests/test_search_db.py
run: uv run pytest cognee/tests/test_search_db.py -v
run-neo4j-pgvector-postgres-search-tests:
name: Search test for Neo4j/PGVector/Postgres (Python ${{ matrix.python-version }})
@ -225,4 +225,4 @@ jobs:
DB_PORT: 5432
DB_USERNAME: cognee
DB_PASSWORD: cognee
run: uv run python ./cognee/tests/test_search_db.py
run: uv run pytest cognee/tests/test_search_db.py -v

View file

@ -1,5 +1,6 @@
import pathlib
import os
import pytest
import cognee
from cognee.infrastructure.databases.graph import get_graph_engine
from cognee.infrastructure.databases.vector import get_vector_engine
@ -22,7 +23,8 @@ from collections import Counter
logger = get_logger()
async def main():
@pytest.mark.asyncio
async def test_search_db():
# This test runs for multiple db settings, to run this locally set the corresponding db envs
await cognee.prune.prune_data()
await cognee.prune.prune_system(metadata=True)
@ -265,7 +267,7 @@ async def main():
await cognee.add(text_1, dataset_name)
await cognee.add([text], dataset_name)
await cognee.add([explanation_file_path_quantum], dataset_name)
await cognee.cognify([dataset_name])
@ -296,9 +298,3 @@ async def main():
assert properties["feedback_weight"] >= 6, (
"Feedback weight calculation is not correct, it should be more then 6."
)
if __name__ == "__main__":
import asyncio
asyncio.run(main())