From ac3f9882787367aeb78306182e24fd0c0f0607f0 Mon Sep 17 00:00:00 2001 From: Igor Ilic Date: Thu, 21 Nov 2024 16:09:50 +0100 Subject: [PATCH] test: Add checking of Networkx database deletion Add test to check if networkx database has been cleaned properly Test COG-488 --- cognee/tests/test_library.py | 12 ++++++++++-- cognee/tests/test_pgvector.py | 2 +- 2 files changed, 11 insertions(+), 3 deletions(-) diff --git a/cognee/tests/test_library.py b/cognee/tests/test_library.py index 9a49206a1..6c9d41800 100755 --- a/cognee/tests/test_library.py +++ b/cognee/tests/test_library.py @@ -57,15 +57,23 @@ async def main(): assert len(history) == 6, "Search history is not correct." + # Assert local data files are cleaned properly await cognee.prune.prune_data() assert not os.path.isdir(data_directory_path), "Local data files are not deleted" + # Assert relational, vector and graph databases have been cleaned properly await cognee.prune.prune_system(metadata=True) + connection = await vector_engine.get_connection() collection_names = await connection.table_names() - assert len(collection_names) == 0, "The vector database is not empty" + assert len(collection_names) == 0, "LanceDB vector database is not empty" + from cognee.infrastructure.databases.relational import get_relational_engine - assert not os.path.exists(get_relational_engine().db_path), "The relational database is not empty" + assert not os.path.exists(get_relational_engine().db_path), "SQLite relational database is not empty" + + from cognee.infrastructure.databases.graph import get_graph_config + graph_config = get_graph_config() + assert not os.path.exists(graph_config.graph_file_path), "Networkx graph database is not empty" if __name__ == "__main__": import asyncio diff --git a/cognee/tests/test_pgvector.py b/cognee/tests/test_pgvector.py index c9c841d31..bd6584cbc 100644 --- a/cognee/tests/test_pgvector.py +++ b/cognee/tests/test_pgvector.py @@ -94,7 +94,7 @@ async def main(): await cognee.prune.prune_system(metadata=True) tables_in_database = await vector_engine.get_table_names() - assert len(tables_in_database) == 0, "The database is not empty" + assert len(tables_in_database) == 0, "PostgreSQL database is not empty" if __name__ == "__main__": import asyncio