test llama_vector_store_factory

xushilundao · Jul 2, 2023 · 5af2f56 · 5af2f56
1 parent 501bbff
commit 5af2f56
Show file tree

Hide file tree

Showing 3 changed files with 33 additions and 4 deletions.
diff --git a/superagi/resource_manager/resource_manager.py b/superagi/resource_manager/resource_manager.py
@@ -13,7 +13,7 @@ class ResourceManager:
     def __init__(self, agent_id: str = None):
         self.agent_id = agent_id
 
-    async def create_llama_document(self, file_path: str):
+    def create_llama_document(self, file_path: str):
         """
         Creates a document index from a given file path.
         """

diff --git a/superagi/worker.py b/superagi/worker.py
@@ -33,9 +33,6 @@ def summarize_resource(agent_id: int, resource_id: int):
     from superagi.types.storage_types import StorageTypes
     from superagi.models.resource import Resource
     from superagi.resource_manager.resource_manager import ResourceManager
-    import boto3
-    import asyncio
-    from io import BytesIO
 
     engine = connect_db()
     Session = sessionmaker(bind=engine)

diff --git a/tests/unit_tests/resource_manager/test_llama_vector_store_factory.py b/tests/unit_tests/resource_manager/test_llama_vector_store_factory.py
@@ -0,0 +1,32 @@
+import pytest
+from unittest.mock import patch
+
+from llama_index.vector_stores import PineconeVectorStore, RedisVectorStore
+
+from superagi.resource_manager.llama_vector_store_factory import LlamaVectorStoreFactory
+from superagi.types.vector_store_types import VectorStoreType
+
+
+def test_llama_vector_store_factory():
+    # Mocking method arguments
+    vector_store_name = VectorStoreType.PINECONE
+    index_name = "test_index_name"
+    factory = LlamaVectorStoreFactory(vector_store_name, index_name)
+
+    # Test case for VectorStoreType.PINECONE
+    with patch.object(PineconeVectorStore, "__init__", return_value=None):
+        vector_store = factory.get_vector_store()
+        assert isinstance(vector_store, PineconeVectorStore)
+
+    # Test case for VectorStoreType.REDIS
+    factory.vector_store_name = VectorStoreType.REDIS
+    with patch.object(RedisVectorStore, "__init__", return_value=None), \
+            patch('superagi.config.config.get_config', return_value=None):
+        vector_store = factory.get_vector_store()
+        assert isinstance(vector_store, RedisVectorStore)
+
+    # Test case for unknown VectorStoreType
+    factory.vector_store_name = "unknown"
+    with pytest.raises(ValueError) as exc_info:
+        factory.get_vector_store()
+    assert str(exc_info.value) == "unknown vector store is not supported yet."