From 4035302dd45e03f8e6f9e318e5ecd98501537e45 Mon Sep 17 00:00:00 2001
From: hajdul88 <52442977+hajdul88@users.noreply.github.com>
Date: Tue, 26 Nov 2024 16:48:09 +0100
Subject: [PATCH] feat: Adds tests for pgvector, qdrant and weaviate

---
 cognee/tests/test_pgvector.py | 4 ++++
 cognee/tests/test_qdrant.py   | 4 ++++
 cognee/tests/test_weaviate.py | 4 ++++
 3 files changed, 12 insertions(+)

diff --git a/cognee/tests/test_pgvector.py b/cognee/tests/test_pgvector.py
index bd6584cbc..1a62fa4d4 100644
--- a/cognee/tests/test_pgvector.py
+++ b/cognee/tests/test_pgvector.py
@@ -3,6 +3,7 @@ import logging
 import pathlib
 import cognee
 from cognee.api.v1.search import SearchType
+from cognee.modules.retrieval.brute_force_triplet_search import brute_force_triplet_search
 
 logging.basicConfig(level=logging.DEBUG)
 
@@ -89,6 +90,9 @@ async def main():
     history = await cognee.get_search_history()
     assert len(history) == 6, "Search history is not correct."
 
+    results = await brute_force_triplet_search('Who has the most experience with graphic design?')
+    assert len(results) > 0
+
     await cognee.prune.prune_data()
     assert not os.path.isdir(data_directory_path), "Local data files are not deleted"
 
diff --git a/cognee/tests/test_qdrant.py b/cognee/tests/test_qdrant.py
index 4c2462c3b..e0df0e980 100644
--- a/cognee/tests/test_qdrant.py
+++ b/cognee/tests/test_qdrant.py
@@ -5,6 +5,7 @@ import logging
 import pathlib
 import cognee
 from cognee.api.v1.search import SearchType
+from cognee.modules.retrieval.brute_force_triplet_search import brute_force_triplet_search
 
 logging.basicConfig(level=logging.DEBUG)
 
@@ -61,6 +62,9 @@ async def main():
     history = await cognee.get_search_history()
     assert len(history) == 6, "Search history is not correct."
 
+    results = await brute_force_triplet_search('Who has the most experience with graphic design?')
+    assert len(results) > 0
+
     await cognee.prune.prune_data()
     assert not os.path.isdir(data_directory_path), "Local data files are not deleted"
 
diff --git a/cognee/tests/test_weaviate.py b/cognee/tests/test_weaviate.py
index c352df13e..726a8ebc2 100644
--- a/cognee/tests/test_weaviate.py
+++ b/cognee/tests/test_weaviate.py
@@ -3,6 +3,7 @@ import logging
 import pathlib
 import cognee
 from cognee.api.v1.search import SearchType
+from cognee.modules.retrieval.brute_force_triplet_search import brute_force_triplet_search
 
 logging.basicConfig(level=logging.DEBUG)
 
@@ -59,6 +60,9 @@ async def main():
     history = await cognee.get_search_history()
     assert len(history) == 6, "Search history is not correct."
 
+    results = await brute_force_triplet_search('Who has the most experience with graphic design?')
+    assert len(results) > 0
+
     await cognee.prune.prune_data()
     assert not os.path.isdir(data_directory_path), "Local data files are not deleted"