From b7df2dacc745ceaa6843a100d15ad95081ee246c Mon Sep 17 00:00:00 2001
From: Vasilije <8619304+Vasilije1990@users.noreply.github.com>
Date: Sat, 1 Jun 2024 21:43:29 +0200
Subject: [PATCH] updated tests

---
 cognee/tests/test_library.py | 34 +++++++++++++++++++++++++---------
 1 file changed, 25 insertions(+), 9 deletions(-)

diff --git a/cognee/tests/test_library.py b/cognee/tests/test_library.py
index 3d1a1b922..a0256a9aa 100755
--- a/cognee/tests/test_library.py
+++ b/cognee/tests/test_library.py
@@ -45,6 +45,31 @@ async def  main():
     Some notable LLMs are OpenAI's GPT series of models (e.g., GPT-3.5 and GPT-4, used in ChatGPT and Microsoft Copilot), Google's PaLM and Gemini (the latter of which is currently used in the chatbot of the same name), xAI's Grok, Meta's LLaMA family of open-source models, Anthropic's Claude models, Mistral AI's open source models, and Databricks' open source DBRX.
     """
 
+
+    await cognee.cognify([ "cs_explanations"])
+
+    search_results = cognee.search("SIMILARITY", "computer science")
+    assert len(search_results) != 0, "The search results list is empty."
+    print("The search results list is not empty.")
+
+    search_results = cognee.search("CATEGORIES", "DefaultGraphModel__default_user")
+    assert len(search_results) != 0, "The search results list is empty."
+    print("The search results list is not empty.")
+
+
+    search_results = cognee.search("NEIGHBOR", "DefaultGraphModel__default_user")
+    assert len(search_results) != 0, "The search results list is empty."
+    print("The search results list is not empty.")
+
+
+    search_results = cognee.search("SUMMARY", "Work and computers")
+    assert len(search_results) != 0, "The search results list is empty."
+    print("The search results list is not empty.")
+
+    search_results = cognee.search("ADJACENT", "DefaultGraphModel__default_user")
+    assert len(search_results) != 0, "The search results list is empty."
+    print("The search results list is not empty.")
+
     #
     # dataset_name = "cs_explanations"
     # await cognee.add(
@@ -55,15 +80,6 @@ async def  main():
     #     dataset_name
     # )
 
-    await cognee.cognify([ "cs_explanations"])
-
-    search_results = cognee.search("SIMILARITY", "computer science")
-    if len(search_results) == 0:
-        print("The similarity search results list is empty.")
-        raise Exception("The search results list is empty.")
-    else:
-        print("The search results list is not empty.")
-
     # search_results = cognee.search("SIMILARITY", "computer science")
     # if len(search_results) == 0:
     #     print("The similarity search results list is empty.")