Spaces:
Runtime error
Runtime error
import pytest | |
from unittest.mock import Mock, patch | |
from swarms.tools.agent_tools import * | |
from swarms.boss.boss_node import BossNodeInitializer, BossNode | |
# For initializing BossNodeInitializer in multiple tests | |
def mock_boss_node_initializer(): | |
llm = Mock() | |
vectorstore = Mock() | |
agent_executor = Mock() | |
max_iterations = 5 | |
boss_node_initializer = BossNodeInitializer(llm, vectorstore, agent_executor, max_iterations) | |
return boss_node_initializer | |
# Test BossNodeInitializer class __init__ method | |
def test_boss_node_initializer_init(mock_boss_node_initializer): | |
with patch('swarms.tools.agent_tools.BabyAGI.from_llm') as mock_from_llm: | |
assert isinstance(mock_boss_node_initializer, BossNodeInitializer) | |
mock_from_llm.assert_called_once() | |
# Test initialize_vectorstore method of BossNodeInitializer class | |
def test_boss_node_initializer_initialize_vectorstore(mock_boss_node_initializer): | |
with patch('swarms.tools.agent_tools.OpenAIEmbeddings') as mock_embeddings, \ | |
patch('swarms.tools.agent_tools.FAISS') as mock_faiss: | |
result = mock_boss_node_initializer.initialize_vectorstore() | |
mock_embeddings.assert_called_once() | |
mock_faiss.assert_called_once() | |
assert result is not None | |
# Test initialize_llm method of BossNodeInitializer class | |
def test_boss_node_initializer_initialize_llm(mock_boss_node_initializer): | |
with patch('swarms.tools.agent_tools.OpenAI') as mock_llm: | |
result = mock_boss_node_initializer.initialize_llm(mock_llm) | |
mock_llm.assert_called_once() | |
assert result is not None | |
# Test create_task method of BossNodeInitializer class | |
def test_boss_node_initializer_create_task(objective, mock_boss_node_initializer): | |
if objective == '': | |
with pytest.raises(ValueError): | |
mock_boss_node_initializer.create_task(objective) | |
else: | |
assert mock_boss_node_initializer.create_task(objective) == {"objective": objective} | |
# Test run method of BossNodeInitializer class | |
def test_boss_node_initializer_run(task, mock_boss_node_initializer): | |
with patch.object(mock_boss_node_initializer, 'baby_agi'): | |
if task == '': | |
with pytest.raises(ValueError): | |
mock_boss_node_initializer.run(task) | |
else: | |
try: | |
mock_boss_node_initializer.run(task) | |
mock_boss_node_initializer.baby_agi.assert_called_once_with(task) | |
except Exception: | |
pytest.fail("Unexpected Error!") | |
# Test BossNode function | |
def test_boss_node(api_key, objective, llm_class, max_iterations): | |
with patch('os.getenv') as mock_getenv, \ | |
patch('swarms.tools.agent_tools.PromptTemplate.from_template') as mock_from_template, \ | |
patch('swarms.tools.agent_tools.LLMChain') as mock_llm_chain, \ | |
patch('swarms.tools.agent_tools.ZeroShotAgent.create_prompt') as mock_create_prompt, \ | |
patch('swarms.tools.agent_tools.ZeroShotAgent') as mock_zero_shot_agent, \ | |
patch('swarms.tools.agent_tools.AgentExecutor.from_agent_and_tools') as mock_from_agent_and_tools, \ | |
patch('swarms.tools.agent_tools.BossNodeInitializer') as mock_boss_node_initializer, \ | |
patch.object(mock_boss_node_initializer, 'create_task') as mock_create_task, \ | |
patch.object(mock_boss_node_initializer, 'run') as mock_run: | |
if api_key == '' or objective == '' or llm_class == '' or max_iterations <= 0: | |
with pytest.raises(ValueError): | |
BossNode(objective, api_key, vectorstore=None, worker_node=None, llm_class=llm_class, max_iterations=max_iterations, verbose=False) | |
else: | |
mock_getenv.return_value = 'valid_key' | |
BossNode(objective, api_key, vectorstore=None, worker_node=None, llm_class=llm_class, max_iterations=max_iterations, verbose=False) | |
mock_from_template.assert_called_once() | |
mock_llm_chain.assert_called_once() | |
mock_create_prompt.assert_called_once() | |
mock_zero_shot_agent.assert_called_once() | |
mock_from_agent_and_tools.assert_called_once() | |
mock_boss_node_initializer.assert_called_once() | |
mock_create_task.assert_called_once() | |
mock_run.assert_called_once() | |