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() | |