agency-swarm-cn-git/tests/test_utils_modules/test_create_agent_template.py

"""Tests for create_agent_template utility."""

from pathlib import Path
from unittest.mock import patch

import pytest

from agency_swarm.utils.create_agent_template import create_agent_template


class TestCreateAgentTemplate:
    """Test the create_agent_template function."""

    def test_create_basic_agent_template(self, tmp_path: Path) -> None:
        """Test creating a basic agent template with minimal parameters."""
        agent_name = "Test Agent"

        result = create_agent_template(agent_name=agent_name, path=str(tmp_path))
        assert result is True

        # Check folder structure (now uses underscores)
        agent_folder = tmp_path / "test_agent"
        assert agent_folder.exists()
        assert agent_folder.is_dir()

        # Check required files
        assert (agent_folder / "__init__.py").exists()
        assert (agent_folder / "test_agent.py").exists()
        assert (agent_folder / "instructions.md").exists()

        # Check required directories
        assert (agent_folder / "tools").exists()
        assert (agent_folder / "files").exists()
        assert (agent_folder / "tools" / "__init__.py").exists()
        assert (agent_folder / "tools" / "ExampleTool.py").exists()

    def test_create_agent_template_without_description(self, tmp_path: Path) -> None:
        """Test creating agent template without description - should use default placeholder."""
        agent_name = "Simple Agent"

        assert create_agent_template(agent_name=agent_name, path=str(tmp_path)) is True

        agent_folder = tmp_path / "simple_agent"
        agent_file = agent_folder / "simple_agent.py"
        agent_content = agent_file.read_text(encoding="utf-8")

        # Should NOT have description line
        assert "description=" not in agent_content
        assert f'name="{agent_name}"' in agent_content
        assert 'instructions="./instructions.md"' in agent_content

        # Instructions file should have structured template with placeholder
        instructions_file = agent_folder / "instructions.md"
        instructions_content = instructions_file.read_text(encoding="utf-8")
        assert "# Role" in instructions_content
        assert "**[insert role" in instructions_content
        assert "# Goals" in instructions_content
        assert "# Process" in instructions_content
        assert "# Output Format" in instructions_content
        assert "# Additional Notes" in instructions_content

    def test_agent_file_content_structure(self, tmp_path: Path) -> None:
        """Test that generated agent file has correct v1.x structure."""
        agent_name = "Data Processor"
        agent_description = "Processes and analyzes data"

        assert (
            create_agent_template(
                agent_name=agent_name,
                agent_description=agent_description,
                path=str(tmp_path),
            )
            is True
        )

        agent_folder = tmp_path / "data_processor"
        agent_file = agent_folder / "data_processor.py"
        agent_content = agent_file.read_text(encoding="utf-8")

        # Check v1.x imports and structure
        assert "from agency_swarm import Agent, ModelSettings" in agent_content
        assert "data_processor = Agent(" in agent_content
        assert f'name="{agent_name}"' in agent_content
        assert f'description="{agent_description}"' in agent_content
        assert 'instructions="./instructions.md"' in agent_content
        assert 'files_folder="./files"' in agent_content
        assert 'tools_folder="./tools"' in agent_content
        assert 'model="gpt-5.4"' in agent_content  # Updated default model
        assert "ModelSettings(" in agent_content
        # gpt-5.4 is a reasoning model - no temperature
        assert "temperature=" not in agent_content

    def test_instructions_file_content(self, tmp_path: Path) -> None:
        """Test that instructions file has correct structured template content."""
        agent_name = "Research Assistant"
        agent_description = "Helps with research tasks"

        assert (
            create_agent_template(
                agent_name=agent_name,
                agent_description=agent_description,
                path=str(tmp_path),
            )
            is True
        )

        agent_folder = tmp_path / "research_assistant"
        instructions_file = agent_folder / "instructions.md"
        instructions_content = instructions_file.read_text(encoding="utf-8")

        # Should have structured template sections
        assert "# Role" in instructions_content
        assert "You are" in instructions_content
        assert agent_description in instructions_content
        assert "# Goals" in instructions_content
        assert "# Process" in instructions_content
        assert "# Output Format" in instructions_content
        assert "# Additional Notes" in instructions_content

    def test_custom_instructions(self, tmp_path: Path) -> None:
        """Test creating agent with custom instructions."""
        agent_name = "Custom Agent"
        agent_description = "Custom description"
        custom_instructions = "These are custom instructions for the agent."

        assert (
            create_agent_template(
                agent_name=agent_name,
                agent_description=agent_description,
                path=str(tmp_path),
                instructions=custom_instructions,
            )
            is True
        )

        agent_folder = tmp_path / "custom_agent"
        instructions_file = agent_folder / "instructions.md"
        instructions_content = instructions_file.read_text(encoding="utf-8")

        assert instructions_content == custom_instructions

    def test_custom_instructions_with_description(self, tmp_path: Path) -> None:
        """Test creating agent with custom instructions and description."""
        agent_name = "Custom Agent"
        agent_description = "Custom description"
        custom_instructions = "You are a specialized agent. Follow these rules: 1. Be precise 2. Be helpful"

        assert (
            create_agent_template(
                agent_name=agent_name,
                agent_description=agent_description,
                instructions=custom_instructions,
                path=str(tmp_path),
            )
            is True
        )

        agent_folder = tmp_path / "custom_agent"

        # Agent file should have description
        agent_file = agent_folder / "custom_agent.py"
        agent_content = agent_file.read_text(encoding="utf-8")
        assert f'description="{agent_description}"' in agent_content

        # Instructions file should have custom content (not auto-generated)
        instructions_file = agent_folder / "instructions.md"
        instructions_content = instructions_file.read_text(encoding="utf-8")
        assert instructions_content == custom_instructions

    def test_custom_instructions_without_description(self, tmp_path: Path) -> None:
        """Test creating agent with custom instructions but no description."""
        agent_name = "Instruction Agent"
        custom_instructions = "You are a specialized agent with custom instructions."

        assert (
            create_agent_template(agent_name=agent_name, instructions=custom_instructions, path=str(tmp_path)) is True
        )

        agent_folder = tmp_path / "instruction_agent"

        # Agent file should NOT have description
        agent_file = agent_folder / "instruction_agent.py"
        agent_content = agent_file.read_text(encoding="utf-8")
        assert "description=" not in agent_content

        # Instructions file should have custom content
        instructions_file = agent_folder / "instructions.md"
        instructions_content = instructions_file.read_text(encoding="utf-8")
        assert instructions_content == custom_instructions

    def test_use_txt_extension(self, tmp_path: Path) -> None:
        """Test creating agent with .txt instructions file."""
        agent_name = "Text Agent"
        agent_description = "Uses txt files"

        assert (
            create_agent_template(
                agent_name=agent_name,
                agent_description=agent_description,
                path=str(tmp_path),
                use_txt=True,
            )
            is True
        )

        agent_folder = tmp_path / "text_agent"

        # Should create .txt file instead of .md
        assert (agent_folder / "instructions.txt").exists()
        assert not (agent_folder / "instructions.md").exists()

        # Agent file should reference .txt extension
        agent_file = agent_folder / "text_agent.py"
        agent_content = agent_file.read_text(encoding="utf-8")
        assert 'instructions="./instructions.txt"' in agent_content

    def test_without_example_tool(self, tmp_path: Path) -> None:
        """Test creating agent without example tool."""
        agent_name = "No Tool Agent"
        agent_description = "Agent without example tool"

        assert (
            create_agent_template(
                agent_name=agent_name,
                agent_description=agent_description,
                path=str(tmp_path),
                include_example_tool=False,
            )
            is True
        )

        agent_folder = tmp_path / "no_tool_agent"

        # Should not create ExampleTool.py
        assert not (agent_folder / "tools" / "ExampleTool.py").exists()
        # But should still create tools folder and __init__.py
        assert (agent_folder / "tools").exists()
        assert (agent_folder / "tools" / "__init__.py").exists()

    def test_init_file_content(self, tmp_path: Path) -> None:
        """Test that __init__.py file has correct imports."""
        agent_name = "Import Test Agent"
        agent_description = "Test imports"

        assert (
            create_agent_template(
                agent_name=agent_name,
                agent_description=agent_description,
                path=str(tmp_path),
            )
            is True
        )

        agent_folder = tmp_path / "import_test_agent"
        init_file = agent_folder / "__init__.py"
        init_content = init_file.read_text(encoding="utf-8")

        assert "from .import_test_agent import import_test_agent" in init_content

    def test_example_tool_content(self, tmp_path: Path) -> None:
        """Test that ExampleTool.py has correct v1.x structure."""
        agent_name = "Tool Test Agent"
        agent_description = "Test tool generation"

        assert (
            create_agent_template(
                agent_name=agent_name,
                agent_description=agent_description,
                path=str(tmp_path),
            )
            is True
        )

        agent_folder = tmp_path / "tool_test_agent"
        tool_file = agent_folder / "tools" / "ExampleTool.py"
        tool_content = tool_file.read_text(encoding="utf-8")

        # Check v1.x tool structure
        assert "from agency_swarm.tools import BaseTool" in tool_content
        assert "from pydantic import Field" in tool_content
        assert "from dotenv import load_dotenv" in tool_content
        assert "load_dotenv()" in tool_content
        assert "class ExampleTool(BaseTool):" in tool_content
        assert "def run(self):" in tool_content
        assert 'if __name__ == "__main__":' in tool_content

    def test_folder_already_exists_error(self, tmp_path: Path) -> None:
        """Test that creating agent in existing folder raises error."""
        agent_name = "Existing Agent"
        agent_description = "Test existing folder"

        # Create the folder first
        (tmp_path / "existing_agent").mkdir()

        # Should raise exception
        with pytest.raises(FileExistsError, match="Folder already exists"):
            create_agent_template(agent_name=agent_name, agent_description=agent_description, path=str(tmp_path))

    def test_agent_name_normalization(self, tmp_path: Path) -> None:
        """Test that agent names with spaces are normalized correctly."""
        agent_name = "Complex Agent Name With Spaces"
        agent_description = "Test name normalization"

        assert (
            create_agent_template(
                agent_name=agent_name,
                agent_description=agent_description,
                path=str(tmp_path),
            )
            is True
        )

        # Should create folder with underscores
        agent_folder = tmp_path / "complex_agent_name_with_spaces"
        assert agent_folder.exists()

        # Check that agent variable is correct in files
        agent_file = agent_folder / "complex_agent_name_with_spaces.py"
        agent_content = agent_file.read_text(encoding="utf-8")
        assert "complex_agent_name_with_spaces = Agent(" in agent_content

        init_file = agent_folder / "__init__.py"
        init_content = init_file.read_text(encoding="utf-8")
        assert "from .complex_agent_name_with_spaces import complex_agent_name_with_spaces" in init_content

    def test_agent_name_with_dots_preserves_version_digits(self, tmp_path: Path) -> None:
        """Agent names containing dots should normalize to underscores without dropping digits."""
        agent_name = "Versioned Agent 1.2"

        assert create_agent_template(agent_name=agent_name, path=str(tmp_path)) is True

        agent_folder = tmp_path / "versioned_agent_1_2"
        assert agent_folder.exists()
        assert (agent_folder / "versioned_agent_1_2.py").exists()

        init_content = (agent_folder / "__init__.py").read_text(encoding="utf-8")
        assert "from .versioned_agent_1_2 import versioned_agent_1_2" in init_content

    @patch("builtins.input")
    def test_interactive_input(self, mock_input, tmp_path: Path) -> None:
        """Test interactive input when name not provided."""
        mock_input.return_value = "Interactive Agent"

        assert create_agent_template(path=str(tmp_path)) is True

        # Should create agent with input name
        agent_folder = tmp_path / "interactive_agent"
        assert agent_folder.exists()

        agent_file = agent_folder / "interactive_agent.py"
        agent_content = agent_file.read_text(encoding="utf-8")
        assert 'name="Interactive Agent"' in agent_content
        # Should NOT have description since we don't prompt for it anymore
        assert "description=" not in agent_content