agency-swarm-cn-git/tests/test_agency_modules/test_ui.py

"""
Tests for Agency Swarm visualization functionality.
"""

import tempfile
from pathlib import Path
from unittest.mock import patch

import pytest

from agency_swarm import Agency, Agent
from agency_swarm.ui import HTMLVisualizationGenerator, LayoutAlgorithms


@pytest.fixture
def sample_agency():
    """Create a sample agency for testing visualization."""
    ceo = Agent(name="CEO", instructions="You are the CEO")
    manager = Agent(name="Manager", instructions="You manage projects")
    worker = Agent(name="Worker", instructions="You do the work")

    agency = Agency(ceo, communication_flows=[ceo > manager, manager > worker])
    return agency


@pytest.fixture
def sample_agency_data():
    """Sample agency data structure for testing."""
    return {
        "nodes": [
            {
                "id": "CEO",
                "type": "agent",
                "data": {"label": "CEO", "description": "You are the CEO", "isEntryPoint": True},
                "position": {"x": 100, "y": 100},
            },
            {
                "id": "Manager",
                "type": "agent",
                "data": {"label": "Manager", "description": "You manage projects", "isEntryPoint": False},
                "position": {"x": 200, "y": 200},
            },
            {
                "id": "Worker",
                "type": "agent",
                "data": {"label": "Worker", "description": "You do the work", "isEntryPoint": False},
                "position": {"x": 300, "y": 300},
            },
        ],
        "edges": [
            {"id": "CEO-Manager", "source": "CEO", "target": "Manager", "type": "communication"},
            {"id": "Manager-Worker", "source": "Manager", "target": "Worker", "type": "communication"},
        ],
        "metadata": {"agencyName": "Test Agency", "totalAgents": 3, "totalTools": 0},
    }


class TestLayoutAlgorithms:
    """Test the layout algorithms."""

    def test_hierarchical_layout_basic(self, sample_agency_data):
        """Test basic hierarchical layout functionality."""
        nodes = sample_agency_data["nodes"]
        edges = sample_agency_data["edges"]

        positions = LayoutAlgorithms.hierarchical_layout(nodes, edges, width=800, height=600)

        # Check that all agents got positions
        assert "CEO" in positions
        assert "Manager" in positions
        assert "Worker" in positions

        # Check that positions have x and y coordinates
        for _node_id, pos in positions.items():
            assert "x" in pos
            assert "y" in pos
            assert isinstance(pos["x"], int | float)
            assert isinstance(pos["y"], int | float)

    def test_hierarchical_layout_entry_points_on_top(self, sample_agency_data):
        """Test that entry points are positioned at the top."""
        nodes = sample_agency_data["nodes"]
        edges = sample_agency_data["edges"]

        positions = LayoutAlgorithms.hierarchical_layout(nodes, edges, width=800, height=600)

        ceo_y = positions["CEO"]["y"]
        manager_y = positions["Manager"]["y"]
        worker_y = positions["Worker"]["y"]

        # CEO (entry point) should be at the top
        assert ceo_y < manager_y
        assert manager_y < worker_y

    def test_hierarchical_layout_with_tools(self):
        """Test hierarchical layout with tools included."""
        nodes = [
            {"id": "CEO", "type": "agent", "data": {"label": "CEO", "isEntryPoint": True}},
            {"id": "tool1", "type": "tool", "data": {"label": "Tool1", "parentAgent": "CEO"}},
        ]
        edges = []

        positions = LayoutAlgorithms.hierarchical_layout(nodes, edges, width=800, height=600)

        assert "CEO" in positions
        assert "tool1" in positions

        # Tool should be positioned relative to its parent agent
        assert positions["tool1"]["x"] != positions["CEO"]["x"] or positions["tool1"]["y"] != positions["CEO"]["y"]

    def test_hierarchical_layout_orphaned_tools(self):
        """Test positioning of tools without parent agents."""
        nodes = [
            {"id": "CEO", "type": "agent", "data": {"label": "CEO", "isEntryPoint": True}},
            {
                "id": "orphan_tool",
                "type": "tool",
                "data": {"label": "Orphan Tool"},  # No parentAgent
            },
        ]
        edges = []

        positions = LayoutAlgorithms.hierarchical_layout(nodes, edges, width=800, height=600)

        assert "CEO" in positions
        assert "orphan_tool" in positions

        # Orphaned tool should be positioned at bottom
        assert positions["orphan_tool"]["y"] > positions["CEO"]["y"]

    def test_apply_layout(self, sample_agency_data):
        """Test the apply_layout method."""
        result = LayoutAlgorithms.apply_layout(sample_agency_data)

        # Check that structure is preserved
        assert "nodes" in result
        assert "edges" in result
        assert "metadata" in result

        # Check that positions were updated
        for node in result["nodes"]:
            assert "position" in node
            assert "x" in node["position"]
            assert "y" in node["position"]


class TestHTMLVisualizationGenerator:
    """Test the HTML visualization generator."""

    def test_init(self):
        """Test HTMLVisualizationGenerator initialization."""
        generator = HTMLVisualizationGenerator()
        assert generator.template_dir.exists()
        assert (generator.template_dir / "visualization.html").exists()
        assert (generator.template_dir / "styles.css").exists()
        assert (generator.template_dir / "visualization.js").exists()

    def test_load_template(self):
        """Test template loading."""
        generator = HTMLVisualizationGenerator()

        # Test loading existing template
        html_content = generator._load_template("visualization.html")
        assert isinstance(html_content, str)
        assert len(html_content) > 0
        assert "html" in html_content.lower()

    def test_load_template_not_found(self):
        """Test error handling for missing template."""
        generator = HTMLVisualizationGenerator()

        with pytest.raises(FileNotFoundError):
            generator._load_template("nonexistent.html")

    @patch("webbrowser.open")
    def test_open_in_browser_success(self, mock_webbrowser):
        """Test opening file in browser successfully."""
        generator = HTMLVisualizationGenerator()
        test_path = Path("/test/path.html")

        generator._open_in_browser(test_path)
        mock_webbrowser.assert_called_once_with(f"file://{test_path}")

    @patch("webbrowser.open", side_effect=Exception("Browser error"))
    @patch("builtins.print")
    def test_open_in_browser_error(self, mock_print, mock_webbrowser):
        """Test error handling when browser fails to open."""
        generator = HTMLVisualizationGenerator()
        test_path = Path("/test/path.html")

        generator._open_in_browser(test_path)

        # Should print error messages
        assert mock_print.call_count >= 2

    def test_generate_interactive_html(self, sample_agency_data):
        """Test interactive HTML generation."""
        generator = HTMLVisualizationGenerator()

        with tempfile.NamedTemporaryFile(mode="w", suffix=".html", delete=False) as tmp:
            output_file = tmp.name

        try:
            with patch.object(generator, "_open_in_browser"):
                result_path = generator.generate_interactive_html(
                    agency_data=sample_agency_data, output_file=output_file, open_browser=False
                )

            assert result_path == str(Path(output_file).resolve())
            assert Path(output_file).exists()

            # Check file content
            with open(output_file) as f:
                content = f.read()

            assert "Test Agency" in content
            assert "CEO" in content
            assert "Manager" in content
            assert "Worker" in content

        finally:
            # Cleanup
            if Path(output_file).exists():
                Path(output_file).unlink()

    def test_generate_component_files(self, sample_agency_data):
        """Test generation of separate component files."""
        generator = HTMLVisualizationGenerator()

        with tempfile.TemporaryDirectory() as tmp_dir:
            files = generator.generate_component_files(agency_data=sample_agency_data, output_dir=tmp_dir)

            # Check that all expected files were created
            assert "html" in files
            assert "css" in files
            assert "js" in files

            for file_path in files.values():
                assert Path(file_path).exists()
                assert Path(file_path).stat().st_size > 0

    @patch.object(HTMLVisualizationGenerator, "generate_interactive_html")
    def test_create_visualization_from_agency(self, mock_generate, sample_agency):
        """Test creating visualization directly from agency."""
        mock_generate.return_value = "/path/to/output.html"

        result = HTMLVisualizationGenerator.create_visualization_from_agency(
            agency=sample_agency,
            output_file="test.html",
            include_tools=True,
            open_browser=False,
        )

        assert result == "/path/to/output.html"
        mock_generate.assert_called_once()


class TestAgencyTUI:
    """Test terminal UI entry points on Agency."""

    def test_tui_delegates_to_visualization(self, sample_agency):
        """Agency.tui should delegate to the visualization entry point."""
        with patch("agency_swarm.agency.visualization.tui") as mock_tui:
            sample_agency.tui(show_reasoning=True, reload=False)

        mock_tui.assert_called_once_with(sample_agency, show_reasoning=True, reload=False)

    def test_terminal_demo_alias_delegates_to_tui(self, sample_agency):
        """Agency.terminal_demo should remain a compatibility alias for tui."""
        with patch.object(sample_agency, "tui") as mock_tui:
            sample_agency.terminal_demo(show_reasoning=True, reload=False)

        mock_tui.assert_called_once_with(show_reasoning=True, reload=False)


class TestAgencyVisualizationIntegration:
    """Test Agency class visualization methods."""

    def test_visualize(self, sample_agency):
        """Test Agency.visualize method."""
        with tempfile.NamedTemporaryFile(mode="w", suffix=".html", delete=False) as tmp:
            output_file = tmp.name

        try:
            with patch("webbrowser.open"):
                result_path = sample_agency.visualize(output_file=output_file, include_tools=True, open_browser=False)

            assert result_path == str(Path(output_file).resolve())
            assert Path(output_file).exists()

            # Check that file contains expected content
            with open(output_file) as f:
                content = f.read()

            assert "CEO" in content
            assert "Manager" in content
            assert "Worker" in content

        finally:
            if Path(output_file).exists():
                Path(output_file).unlink()

    def test_visualize_import_error(self, sample_agency):
        """Test handling of import errors in visualization."""
        # Patch the import inside the visualize method
        with patch("agency_swarm.agency.Agency.visualize") as mock_method:
            mock_method.side_effect = ImportError("Visualization module not available")
            with pytest.raises(ImportError, match="Visualization module not available"):
                sample_agency.visualize(open_browser=False)

    def test_get_agency_graph_basic(self, sample_agency):
        """Test basic agency graph generation."""
        structure = sample_agency.get_agency_graph()

        assert "nodes" in structure
        assert "edges" in structure
        assert "metadata" not in structure

        # Check nodes
        nodes = structure["nodes"]
        assert len(nodes) >= 3  # At least CEO, Manager, Worker

        agent_nodes = [n for n in nodes if n["type"] == "agent"]
        assert len(agent_nodes) == 3

        # Check edges
        edges = structure["edges"]
        communication_edges = [e for e in edges if e["type"] == "communication"]
        assert len(communication_edges) >= 2  # CEO->Manager, Manager->Worker

    def test_get_agency_graph_with_tools(self, sample_agency):
        """Test agency graph generation with tools included."""
        # Test that the method works with include_tools=True
        # We'll test the structure without actually adding tools to avoid tool type complications
        structure = sample_agency.get_agency_graph(include_tools=True)

        assert "nodes" in structure
        assert "edges" in structure
        assert "metadata" not in structure

        # Check that the structure is valid even when no tools are present
        agent_nodes = [n for n in structure["nodes"] if n["type"] == "agent"]
        assert len(agent_nodes) >= 3

        # Tool nodes and edges should be empty if no tools are added, which is fine
        tool_nodes = [n for n in structure["nodes"] if n["type"] == "tool"]
        tool_edges = [e for e in structure["edges"] if e["type"] == "tool"]

        # These should be lists, even if empty
        assert isinstance(tool_nodes, list)
        assert isinstance(tool_edges, list)

    def test_get_agency_graph_without_tools(self, sample_agency):
        """Test agency graph generation without tools."""
        structure = sample_agency.get_agency_graph(include_tools=False)

        # Should only have agent nodes
        tool_nodes = [n for n in structure["nodes"] if n["type"] == "tool"]
        assert len(tool_nodes) == 0

        # Should only have communication edges
        tool_edges = [e for e in structure["edges"] if e["type"] == "tool"]
        assert len(tool_edges) == 0

    def test_get_agency_graph_hierarchical_layout(self, sample_agency):
        """Test hierarchical layout in get_agency_graph."""
        # Test hierarchical layout
        structure = sample_agency.get_agency_graph()

        # Check that nodes have positions
        for node in structure["nodes"]:
            assert "position" in node
            assert "x" in node["position"]
            assert "y" in node["position"]

    def test_layout_algorithms_manager_vs_leaf_positioning(self):
        """Test that manager agents and leaf agents position tools differently."""
        # Create a more complex structure with manager and leaf agents
        nodes = [
            {"id": "Manager", "type": "agent", "data": {"label": "Manager", "isEntryPoint": True}},
            {"id": "Worker1", "type": "agent", "data": {"label": "Worker1", "isEntryPoint": False}},
            {"id": "Worker2", "type": "agent", "data": {"label": "Worker2", "isEntryPoint": False}},
            {"id": "manager_tool", "type": "tool", "data": {"label": "Manager Tool", "parentAgent": "Manager"}},
            {"id": "worker_tool", "type": "tool", "data": {"label": "Worker Tool", "parentAgent": "Worker1"}},
        ]

        # Manager has multiple outgoing connections
        edges = [
            {"source": "Manager", "target": "Worker1", "type": "communication"},
            {"source": "Manager", "target": "Worker2", "type": "communication"},
        ]

        positions = LayoutAlgorithms.hierarchical_layout(nodes, edges)

        # Both tools should be positioned
        assert "manager_tool" in positions
        assert "worker_tool" in positions

        # Manager tool should be positioned to the right (x > manager x)
        # Worker tool should be positioned below (y > worker y)
        manager_pos = positions["Manager"]
        worker_pos = positions["Worker1"]
        manager_tool_pos = positions["manager_tool"]
        worker_tool_pos = positions["worker_tool"]

        # These assertions test the smart positioning logic
        assert manager_tool_pos["x"] > manager_pos["x"]  # Tool to the right of manager
        assert worker_tool_pos["y"] > worker_pos["y"]  # Tool below worker