tests/swarms.py

import pytest
import logging
from unittest.mock import patch
from swarms.swarms.swarms import HierarchicalSwarm  # replace with your actual module name

@pytest.fixture
def swarm():
    return HierarchicalSwarm(
        model_id='gpt-4', 
        openai_api_key='some_api_key', 
        use_vectorstore=True, 
        embedding_size=1024, 
        use_async=False, 
        human_in_the_loop=True, 
        model_type='openai', 
        boss_prompt='boss', 
        worker_prompt='worker', 
        temperature=0.5,
        max_iterations=100,
        logging_enabled=True
    )

@pytest.fixture
def swarm_no_logging():
    return HierarchicalSwarm(logging_enabled=False)

def test_swarm_init(swarm):
    assert swarm.model_id == 'gpt-4'
    assert swarm.openai_api_key == 'some_api_key'
    assert swarm.use_vectorstore
    assert swarm.embedding_size == 1024
    assert not swarm.use_async
    assert swarm.human_in_the_loop
    assert swarm.model_type == 'openai'
    assert swarm.boss_prompt == 'boss'
    assert swarm.worker_prompt == 'worker'
    assert swarm.temperature == 0.5
    assert swarm.max_iterations == 100
    assert swarm.logging_enabled
    assert isinstance(swarm.logger, logging.Logger)

def test_swarm_no_logging_init(swarm_no_logging):
    assert not swarm_no_logging.logging_enabled
    assert swarm_no_logging.logger.disabled

@patch('your_module.OpenAI')
@patch('your_module.HuggingFaceLLM')
def test_initialize_llm(mock_huggingface, mock_openai, swarm):
    swarm.initialize_llm('openai')
    mock_openai.assert_called_once_with(openai_api_key='some_api_key', temperature=0.5)
    
    swarm.initialize_llm('huggingface')
    mock_huggingface.assert_called_once_with(model_id='gpt-4', temperature=0.5)

@patch('your_module.HierarchicalSwarm.initialize_llm')
def test_initialize_tools(mock_llm, swarm):
    mock_llm.return_value = 'mock_llm_class'
    tools = swarm.initialize_tools('openai')
    assert 'mock_llm_class' in tools

@patch('your_module.HierarchicalSwarm.initialize_llm')
def test_initialize_tools_with_extra_tools(mock_llm, swarm):
    mock_llm.return_value = 'mock_llm_class'
    tools = swarm.initialize_tools('openai', extra_tools=['tool1', 'tool2'])
    assert 'tool1' in tools
    assert 'tool2' in tools

@patch('your_module.OpenAIEmbeddings')
@patch('your_module.FAISS')
def test_initialize_vectorstore(mock_faiss, mock_openai_embeddings, swarm):
    mock_openai_embeddings.return_value.embed_query = 'embed_query'
    swarm.initialize_vectorstore()
    mock_faiss.assert_called_once_with('embed_query', instance_of(faiss.IndexFlatL2), instance_of(InMemoryDocstore), {})