You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
134 lines
5.0 KiB
134 lines
5.0 KiB
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
|
|
@pytest.fixture
|
|
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
|
|
@pytest.mark.parametrize("objective", ["valid objective", ""])
|
|
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
|
|
@pytest.mark.parametrize("task", ["valid task", ""])
|
|
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
|
|
@pytest.mark.parametrize(
|
|
"api_key, objective, llm_class, max_iterations",
|
|
[
|
|
("valid_key", "valid_objective", OpenAI, 5),
|
|
("", "valid_objective", OpenAI, 5),
|
|
("valid_key", "", OpenAI, 5),
|
|
("valid_key", "valid_objective", "", 5),
|
|
("valid_key", "valid_objective", OpenAI, 0),
|
|
],
|
|
)
|
|
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()
|