test document

TheSnowGuru · Jul 2, 2023 · 0462efd · 0462efd
1 parent 89118d0
commit 0462efd
Show file tree

Hide file tree

Showing 3 changed files with 54 additions and 2 deletions.
diff --git a/superagi/resource_manager/llama_document_summary.py b/superagi/resource_manager/llama_document_summary.py
@@ -13,7 +13,7 @@ def __init__(self, model_name=get_config("RESOURCES_EMBEDDING_MODEL_NAME", "gpt-
     def generate_summary_of_document(self, documents: list[Document]):
         from llama_index import LLMPredictor, ServiceContext, ResponseSynthesizer, DocumentSummaryIndex
 
-        os.environ["OPENAI_API_KEY"] = get_config("OPENAI_API_KEY")
+        os.environ["OPENAI_API_KEY"] = get_config("OPENAI_API_KEY", "")
         llm_predictor_chatgpt = LLMPredictor(llm=self._build_llm())
         service_context = ServiceContext.from_defaults(llm_predictor=llm_predictor_chatgpt, chunk_size=1024)
         response_synthesizer = ResponseSynthesizer.from_args(response_mode=ResponseMode.TREE_SUMMARIZE, use_async=True)

diff --git a/superagi/resource_manager/resource_manager.py b/superagi/resource_manager/resource_manager.py
@@ -55,7 +55,7 @@ def save_document_to_vector_store(self, documents: list, resource_id: str):
         from llama_index import VectorStoreIndex, StorageContext
         import openai
         openai.api_key = get_config("OPENAI_API_KEY")
-        os.environ["OPENAI_API_KEY"] = get_config("OPENAI_API_KEY")
+        os.environ["OPENAI_API_KEY"] = get_config("OPENAI_API_KEY", "")
         for docs in documents:
             if docs.metadata is None:
                 docs.metadata = {}

diff --git a/tests/unit_tests/resource_manager/test_llama_document_creation.py b/tests/unit_tests/resource_manager/test_llama_document_creation.py
@@ -0,0 +1,52 @@
+import pytest
+from unittest.mock import patch, MagicMock
+from superagi.resource_manager.resource_manager import ResourceManager
+
+
+def test_create_llama_document_s3(mocker):
+    agent_id = 'test_agent'
+    resource_manager = ResourceManager(agent_id)
+
+    mock_boto_client = MagicMock()
+    mock_s3_obj = {
+        'Body': MagicMock(read=MagicMock(return_value='mock_file_content'))
+    }
+    mock_boto_client.get_object.return_value = mock_s3_obj
+    mocker.patch('boto3.client', return_value=mock_boto_client)
+
+    mocker.patch('superagi.resource_manager.resource_manager.get_config',
+                 side_effect=['mock_access_key', 'mock_secret_key', 'mock_bucket'])
+
+    mocker.patch('builtins.open', mocker.mock_open())
+    mocker.patch('os.remove')
+
+    MockSimpleDirectoryReader = MagicMock()
+    mocker.patch('superagi.resource_manager.resource_manager.SimpleDirectoryReader',
+                 return_value=MockSimpleDirectoryReader)
+
+    resource_manager.create_llama_document_s3('mock_file_path')
+
+    mock_boto_client.get_object.assert_called_once_with(
+        Bucket='mock_bucket',
+        Key='mock_file_path')
+    MockSimpleDirectoryReader.load_data.assert_called_once()
+
+
+def test_create_llama_document_s3_file_path_provided(mocker):
+    resource_manager = ResourceManager('test_agent')
+
+    mock_boto_client = MagicMock()
+    mocker.patch('boto3.client', return_value=mock_boto_client)
+
+    mocker.patch('superagi.resource_manager.resource_manager.get_config',
+                 side_effect=['mock_access_key', 'mock_secret_key', 'mock_bucket'])
+
+    mocker.patch('builtins.open', mocker.mock_open())
+    mocker.patch('os.remove')
+
+    MockSimpleDirectoryReader = MagicMock()
+    mocker.patch('superagi.resource_manager.resource_manager.SimpleDirectoryReader',
+                 return_value=MockSimpleDirectoryReader)
+
+    with pytest.raises(Exception, match="file_path must be provided"):
+        resource_manager.create_llama_document_s3(None)