Spaces:
Sleeping
Sleeping
File size: 1,694 Bytes
d8d14f1 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 |
import pytest
from swarm_models import OpenAIChat
from swarms.structs.agent import Agent
from swarms.structs.company import Company
# Mock OpenAIChat instance
llm = OpenAIChat(openai_api_key="test_key", max_tokens=4000)
# Mock Agents
ceo = Agent(llm=llm, name="CEO")
dev = Agent(llm=llm, name="Developer")
va = Agent(llm=llm, name="VA")
hr = Agent(llm=llm, name="HR")
shared_instructions = "Listen to your boss"
def test_add_agent():
company = Company(
org_chart=[[ceo, [dev, va]]],
shared_instructions=shared_instructions,
)
company.add(hr)
assert hr in company.agents
def test_get_agent():
company = Company(
org_chart=[[ceo, [dev, va]]],
shared_instructions=shared_instructions,
)
company.add(hr)
assert company.get("HR") == hr
def test_remove_agent():
company = Company(
org_chart=[[ceo, [dev, va]]],
shared_instructions=shared_instructions,
)
company.add(hr)
company.remove(hr)
assert hr not in company.agents
def test_add_existing_agent():
company = Company(
org_chart=[[ceo, [dev, va]]],
shared_instructions=shared_instructions,
)
company.add(hr)
with pytest.raises(ValueError):
company.add(hr)
def test_get_nonexistent_agent():
company = Company(
org_chart=[[ceo, [dev, va]]],
shared_instructions=shared_instructions,
)
with pytest.raises(ValueError):
company.get("Nonexistent")
def test_remove_nonexistent_agent():
company = Company(
org_chart=[[ceo, [dev, va]]],
shared_instructions=shared_instructions,
)
with pytest.raises(ValueError):
company.remove(hr)
|