diff --git a/cognee/tests/test_library.py b/cognee/tests/test_library.py
index e279bcc51..2c656247f 100755
--- a/cognee/tests/test_library.py
+++ b/cognee/tests/test_library.py
@@ -48,25 +48,25 @@ async def  main():
 
     await cognee.cognify([ "cs_explanations"])
 
-    search_results = await cognee.search("SIMILARITY", "computer science")
+    search_results = await cognee.search("SIMILARITY", { "query": "computer science" })
     assert len(search_results) != 0, "The search results list is empty."
     print("The search results list is not empty.")
 
-    search_results = await cognee.search("CATEGORIES", "DefaultGraphModel__default_user")
+    search_results = await cognee.search("CATEGORIES", { "query": "DefaultGraphModel__default_user" })
     assert len(search_results) != 0, "The search results list is empty."
     print("The search results list is not empty.")
 
 
-    search_results = await cognee.search("NEIGHBOR", "DefaultGraphModel__default_user")
+    search_results = await cognee.search("NEIGHBOR", { "query": "DefaultGraphModel__default_user" })
     assert len(search_results) != 0, "The search results list is empty."
     print("The search results list is not empty.")
 
 
-    search_results = await cognee.search("SUMMARY", "Work and computers")
+    search_results = await cognee.search("SUMMARY", { "query": "Work and computers" })
     assert len(search_results) != 0, "The search results list is empty."
     print("The search results list is not empty.")
 
-    search_results = await cognee.search("ADJACENT", "DefaultGraphModel__default_user")
+    search_results = await cognee.search("ADJACENT", { "query": "DefaultGraphModel__default_user" })
     assert len(search_results) != 0, "The search results list is empty."
     print("The search results list is not empty.")