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

converts search_db test to pytest test

Browse source
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
Show stats Download patch file Download diff file Expand all files Collapse all files

8
.github/workflows/search_db_tests.yml vendored
View file

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

12
cognee/tests/test_search_db.py
View file

@ -1,5 +1,6 @@
import pathlib import pathlib
import os import os
import pytest
import cognee import cognee
from cognee.infrastructure.databases.graph import get_graph_engine from cognee.infrastructure.databases.graph import get_graph_engine
from cognee.infrastructure.databases.vector import get_vector_engine from cognee.infrastructure.databases.vector import get_vector_engine
@ -22,7 +23,8 @@ from collections import Counter
logger = get_logger() 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 # 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_data()
await cognee.prune.prune_system(metadata=True) 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_1, dataset_name)
await cognee.add([text], dataset_name) await cognee.add([explanation_file_path_quantum], dataset_name)
await cognee.cognify([dataset_name]) await cognee.cognify([dataset_name])
@ -296,9 +298,3 @@ async def main():
assert properties["feedback_weight"] >= 6, ( assert properties["feedback_weight"] >= 6, (
"Feedback weight calculation is not correct, it should be more then 6." "Feedback weight calculation is not correct, it should be more then 6."
) )
if __name__ == "__main__":
import asyncio
asyncio.run(main())