Cog 488 test database deletion (#16)

Added testing of database deletion for every end to end test

cognee/infrastructure/databases/relational/sqlalchemy/SqlAlchemyAdapter.py

@@ -130,6 +130,29 @@ class SQLAlchemyAdapter():
                     return metadata.tables[full_table_name]
                 raise ValueError(f"Table '{full_table_name}' not found.")
 
+    async def get_table_names(self) -> List[str]:
+        """
+        Return a list of all tables names in database
+        """
+        table_names = []
+        async with self.engine.begin() as connection:
+            if self.engine.dialect.name == "sqlite":
+                await connection.run_sync(Base.metadata.reflect)
+                for table in Base.metadata.tables:
+                    table_names.append(str(table))
+            else:
+                schema_list = await self.get_schema_list()
+                # Create a MetaData instance to load table information
+                metadata = MetaData()
+                # Drop all tables from all schemas
+                for schema_name in schema_list:
+                    # Load the schema information into the MetaData object
+                    await connection.run_sync(metadata.reflect, schema=schema_name)
+                    for table in metadata.sorted_tables:
+                        table_names.append(str(table))
+                    metadata.clear()
+        return table_names
+
 
     async def get_data(self, table_name: str, filters: dict = None):
         async with self.engine.begin() as connection:

cognee/tests/test_library.py

@@ -57,6 +57,24 @@ 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, "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), "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
     asyncio.run(main(), debug=True)

cognee/tests/test_neo4j.py

@@ -61,6 +61,15 @@ async def main():
 
     assert len(history) == 6, "Search history is not correct."
 
+    await cognee.prune.prune_data()
+    assert not os.path.isdir(data_directory_path), "Local data files are not deleted"
+
+    await cognee.prune.prune_system(metadata=True)
+    from cognee.infrastructure.databases.graph import get_graph_engine
+    graph_engine = await get_graph_engine()
+    nodes, edges = await graph_engine.get_graph_data()
+    assert len(nodes) == 0 and len(edges) == 0, "Neo4j graph database is not empty"
+
 if __name__ == "__main__":
     import asyncio
     asyncio.run(main())

cognee/tests/test_pgvector.py

@@ -87,9 +87,15 @@ async def main():
         print(f"{result}\n")
 
     history = await cognee.get_search_history()
-
     assert len(history) == 6, "Search history is not correct."
 
+    await cognee.prune.prune_data()
+    assert not os.path.isdir(data_directory_path), "Local data files are not deleted"
+
+    await cognee.prune.prune_system(metadata=True)
+    tables_in_database = await vector_engine.get_table_names()
+    assert len(tables_in_database) == 0, "PostgreSQL database is not empty"
+
 if __name__ == "__main__":
     import asyncio
 

cognee/tests/test_qdrant.py

@@ -59,9 +59,16 @@ async def main():
         print(f"{result}\n")
 
     history = await cognee.get_search_history()
-
     assert len(history) == 6, "Search history is not correct."
 
+    await cognee.prune.prune_data()
+    assert not os.path.isdir(data_directory_path), "Local data files are not deleted"
+
+    await cognee.prune.prune_system(metadata=True)
+    qdrant_client = get_vector_engine().get_qdrant_client()
+    collections_response = await qdrant_client.get_collections()
+    assert len(collections_response.collections) == 0, "QDrant vector database is not empty"
+
 if __name__ == "__main__":
     import asyncio
     asyncio.run(main())

cognee/tests/test_weaviate.py

@@ -57,9 +57,15 @@ async def main():
         print(f"{result}\n")
 
     history = await cognee.get_search_history()
-
     assert len(history) == 6, "Search history is not correct."
 
+    await cognee.prune.prune_data()
+    assert not os.path.isdir(data_directory_path), "Local data files are not deleted"
+
+    await cognee.prune.prune_system(metadata=True)
+    collections =  get_vector_engine().client.collections.list_all()
+    assert len(collections) == 0, "Weaviate vector database is not empty"
+
 if __name__ == "__main__":
     import asyncio
     asyncio.run(main())