haproxy-mcp/tests/unit/tools/test_servers.py

"""Unit tests for server management tools."""

import json
from unittest.mock import patch, MagicMock

import pytest

from haproxy_mcp.exceptions import HaproxyError


class TestHaproxyListServers:
    """Tests for haproxy_list_servers tool function."""

    def test_list_servers_invalid_domain(self, patch_config_paths):
        """Reject invalid domain format."""
        from haproxy_mcp.tools.servers import register_server_tools
        mcp = MagicMock()
        registered_tools = {}

        def capture_tool():
            def decorator(func):
                registered_tools[func.__name__] = func
                return func
            return decorator

        mcp.tool = capture_tool
        register_server_tools(mcp)

        result = registered_tools["haproxy_list_servers"](domain="-invalid")

        assert "Error" in result
        assert "Invalid domain" in result

    def test_list_servers_empty_backend(self, mock_socket_class, mock_select, patch_config_paths, response_builder):
        """List servers for domain with no servers."""
        with open(patch_config_paths["map_file"], "w") as f:
            f.write("example.com pool_1\n")

        mock_sock = mock_socket_class(responses={
            "show servers state": response_builder.servers_state([
                {"be_name": "pool_1", "srv_name": "pool_1_1", "srv_addr": "0.0.0.0", "srv_port": 0},
            ]),
        })

        with patch("socket.socket", return_value=mock_sock):
            from haproxy_mcp.tools.servers import register_server_tools
            mcp = MagicMock()
            registered_tools = {}

            def capture_tool():
                def decorator(func):
                    registered_tools[func.__name__] = func
                    return func
                return decorator

            mcp.tool = capture_tool
            register_server_tools(mcp)

            result = registered_tools["haproxy_list_servers"](domain="example.com")

        assert "pool_1" in result
        assert "disabled" in result

    def test_list_servers_with_active_servers(self, mock_socket_class, mock_select, patch_config_paths, response_builder):
        """List servers with active servers."""
        with open(patch_config_paths["map_file"], "w") as f:
            f.write("example.com pool_1\n")

        mock_sock = mock_socket_class(responses={
            "show servers state": response_builder.servers_state([
                {"be_name": "pool_1", "srv_name": "pool_1_1", "srv_addr": "10.0.0.1", "srv_port": 80},
                {"be_name": "pool_1", "srv_name": "pool_1_2", "srv_addr": "10.0.0.2", "srv_port": 80},
            ]),
        })

        with patch("socket.socket", return_value=mock_sock):
            from haproxy_mcp.tools.servers import register_server_tools
            mcp = MagicMock()
            registered_tools = {}

            def capture_tool():
                def decorator(func):
                    registered_tools[func.__name__] = func
                    return func
                return decorator

            mcp.tool = capture_tool
            register_server_tools(mcp)

            result = registered_tools["haproxy_list_servers"](domain="example.com")

        assert "10.0.0.1" in result
        assert "10.0.0.2" in result
        assert "active" in result


class TestHaproxyAddServer:
    """Tests for haproxy_add_server tool function."""

    def test_add_server_invalid_domain(self, patch_config_paths):
        """Reject invalid domain format."""
        from haproxy_mcp.tools.servers import register_server_tools
        mcp = MagicMock()
        registered_tools = {}

        def capture_tool():
            def decorator(func):
                registered_tools[func.__name__] = func
                return func
            return decorator

        mcp.tool = capture_tool
        register_server_tools(mcp)

        result = registered_tools["haproxy_add_server"](
            domain="-invalid",
            slot=1,
            ip="10.0.0.1",
            http_port=80
        )

        assert "Error" in result
        assert "Invalid domain" in result

    def test_add_server_empty_ip(self, patch_config_paths):
        """Reject empty IP address."""
        from haproxy_mcp.tools.servers import register_server_tools
        mcp = MagicMock()
        registered_tools = {}

        def capture_tool():
            def decorator(func):
                registered_tools[func.__name__] = func
                return func
            return decorator

        mcp.tool = capture_tool
        register_server_tools(mcp)

        result = registered_tools["haproxy_add_server"](
            domain="example.com",
            slot=1,
            ip="",
            http_port=80
        )

        assert "Error" in result
        assert "IP address is required" in result

    def test_add_server_invalid_ip(self, patch_config_paths):
        """Reject invalid IP address."""
        from haproxy_mcp.tools.servers import register_server_tools
        mcp = MagicMock()
        registered_tools = {}

        def capture_tool():
            def decorator(func):
                registered_tools[func.__name__] = func
                return func
            return decorator

        mcp.tool = capture_tool
        register_server_tools(mcp)

        result = registered_tools["haproxy_add_server"](
            domain="example.com",
            slot=1,
            ip="not-an-ip",
            http_port=80
        )

        assert "Error" in result
        assert "Invalid IP" in result

    def test_add_server_invalid_port(self, patch_config_paths):
        """Reject invalid port."""
        from haproxy_mcp.tools.servers import register_server_tools
        mcp = MagicMock()
        registered_tools = {}

        def capture_tool():
            def decorator(func):
                registered_tools[func.__name__] = func
                return func
            return decorator

        mcp.tool = capture_tool
        register_server_tools(mcp)

        result = registered_tools["haproxy_add_server"](
            domain="example.com",
            slot=1,
            ip="10.0.0.1",
            http_port=70000
        )

        assert "Error" in result
        assert "Port" in result

    def test_add_server_invalid_slot(self, patch_config_paths):
        """Reject invalid slot number."""
        from haproxy_mcp.tools.servers import register_server_tools
        mcp = MagicMock()
        registered_tools = {}

        def capture_tool():
            def decorator(func):
                registered_tools[func.__name__] = func
                return func
            return decorator

        mcp.tool = capture_tool
        register_server_tools(mcp)

        result = registered_tools["haproxy_add_server"](
            domain="example.com",
            slot=99,  # > MAX_SLOTS
            ip="10.0.0.1",
            http_port=80
        )

        assert "Error" in result
        assert "Slot" in result

    def test_add_server_success(self, mock_socket_class, mock_select, patch_config_paths, response_builder):
        """Successfully add server."""
        with open(patch_config_paths["map_file"], "w") as f:
            f.write("example.com pool_1\n")

        mock_sock = mock_socket_class(responses={
            "set server": "",
        })

        with patch("socket.socket", return_value=mock_sock):
            from haproxy_mcp.tools.servers import register_server_tools
            mcp = MagicMock()
            registered_tools = {}

            def capture_tool():
                def decorator(func):
                    registered_tools[func.__name__] = func
                    return func
                return decorator

            mcp.tool = capture_tool
            register_server_tools(mcp)

            result = registered_tools["haproxy_add_server"](
                domain="example.com",
                slot=1,
                ip="10.0.0.1",
                http_port=8080
            )

        assert "example.com" in result
        assert "slot 1" in result
        assert "10.0.0.1:8080" in result

    def test_add_server_auto_slot(self, mock_socket_class, mock_select, patch_config_paths, response_builder):
        """Auto-select slot when slot=0."""
        with open(patch_config_paths["map_file"], "w") as f:
            f.write("example.com pool_1\n")

        mock_sock = mock_socket_class(responses={
            "show servers state": response_builder.servers_state([
                {"be_name": "pool_1", "srv_name": "pool_1_1", "srv_addr": "10.0.0.1", "srv_port": 80},
                {"be_name": "pool_1", "srv_name": "pool_1_2", "srv_addr": "0.0.0.0", "srv_port": 0},
            ]),
            "set server": "",
        })

        with patch("socket.socket", return_value=mock_sock):
            from haproxy_mcp.tools.servers import register_server_tools
            mcp = MagicMock()
            registered_tools = {}

            def capture_tool():
                def decorator(func):
                    registered_tools[func.__name__] = func
                    return func
                return decorator

            mcp.tool = capture_tool
            register_server_tools(mcp)

            result = registered_tools["haproxy_add_server"](
                domain="example.com",
                slot=0,  # Auto-select
                ip="10.0.0.2",
                http_port=80
            )

        assert "slot 2" in result  # First available slot


class TestHaproxyAddServers:
    """Tests for haproxy_add_servers tool function."""

    def test_add_servers_invalid_domain(self, patch_config_paths):
        """Reject invalid domain format."""
        from haproxy_mcp.tools.servers import register_server_tools
        mcp = MagicMock()
        registered_tools = {}

        def capture_tool():
            def decorator(func):
                registered_tools[func.__name__] = func
                return func
            return decorator

        mcp.tool = capture_tool
        register_server_tools(mcp)

        result = registered_tools["haproxy_add_servers"](
            domain="-invalid",
            servers='[{"slot": 1, "ip": "10.0.0.1"}]'
        )

        assert "Error" in result
        assert "Invalid domain" in result

    def test_add_servers_invalid_json(self, patch_config_paths):
        """Reject invalid JSON."""
        from haproxy_mcp.tools.servers import register_server_tools
        mcp = MagicMock()
        registered_tools = {}

        def capture_tool():
            def decorator(func):
                registered_tools[func.__name__] = func
                return func
            return decorator

        mcp.tool = capture_tool
        register_server_tools(mcp)

        result = registered_tools["haproxy_add_servers"](
            domain="example.com",
            servers='not valid json'
        )

        assert "Error" in result
        assert "Invalid JSON" in result

    def test_add_servers_not_array(self, patch_config_paths):
        """Reject non-array JSON."""
        from haproxy_mcp.tools.servers import register_server_tools
        mcp = MagicMock()
        registered_tools = {}

        def capture_tool():
            def decorator(func):
                registered_tools[func.__name__] = func
                return func
            return decorator

        mcp.tool = capture_tool
        register_server_tools(mcp)

        result = registered_tools["haproxy_add_servers"](
            domain="example.com",
            servers='{"slot": 1, "ip": "10.0.0.1"}'
        )

        assert "Error" in result
        assert "must be a JSON array" in result

    def test_add_servers_empty_array(self, patch_config_paths):
        """Reject empty array."""
        from haproxy_mcp.tools.servers import register_server_tools
        mcp = MagicMock()
        registered_tools = {}

        def capture_tool():
            def decorator(func):
                registered_tools[func.__name__] = func
                return func
            return decorator

        mcp.tool = capture_tool
        register_server_tools(mcp)

        result = registered_tools["haproxy_add_servers"](
            domain="example.com",
            servers='[]'
        )

        assert "Error" in result
        assert "empty" in result

    def test_add_servers_duplicate_slots(self, patch_config_paths):
        """Reject duplicate slot numbers."""
        from haproxy_mcp.tools.servers import register_server_tools
        mcp = MagicMock()
        registered_tools = {}

        def capture_tool():
            def decorator(func):
                registered_tools[func.__name__] = func
                return func
            return decorator

        mcp.tool = capture_tool
        register_server_tools(mcp)

        result = registered_tools["haproxy_add_servers"](
            domain="example.com",
            servers='[{"slot": 1, "ip": "10.0.0.1"}, {"slot": 1, "ip": "10.0.0.2"}]'
        )

        assert "Error" in result
        assert "Duplicate" in result

    def test_add_servers_success(self, mock_socket_class, mock_select, patch_config_paths, response_builder):
        """Successfully add multiple servers."""
        with open(patch_config_paths["map_file"], "w") as f:
            f.write("example.com pool_1\n")

        mock_sock = mock_socket_class(responses={
            "set server": "",
        })

        with patch("socket.socket", return_value=mock_sock):
            from haproxy_mcp.tools.servers import register_server_tools
            mcp = MagicMock()
            registered_tools = {}

            def capture_tool():
                def decorator(func):
                    registered_tools[func.__name__] = func
                    return func
                return decorator

            mcp.tool = capture_tool
            register_server_tools(mcp)

            result = registered_tools["haproxy_add_servers"](
                domain="example.com",
                servers='[{"slot": 1, "ip": "10.0.0.1"}, {"slot": 2, "ip": "10.0.0.2"}]'
            )

        assert "Added 2 servers" in result
        assert "slot 1" in result
        assert "slot 2" in result


class TestHaproxyRemoveServer:
    """Tests for haproxy_remove_server tool function."""

    def test_remove_server_invalid_domain(self, patch_config_paths):
        """Reject invalid domain format."""
        from haproxy_mcp.tools.servers import register_server_tools
        mcp = MagicMock()
        registered_tools = {}

        def capture_tool():
            def decorator(func):
                registered_tools[func.__name__] = func
                return func
            return decorator

        mcp.tool = capture_tool
        register_server_tools(mcp)

        result = registered_tools["haproxy_remove_server"](
            domain="-invalid",
            slot=1
        )

        assert "Error" in result
        assert "Invalid domain" in result

    def test_remove_server_invalid_slot(self, patch_config_paths):
        """Reject invalid slot number."""
        from haproxy_mcp.tools.servers import register_server_tools
        mcp = MagicMock()
        registered_tools = {}

        def capture_tool():
            def decorator(func):
                registered_tools[func.__name__] = func
                return func
            return decorator

        mcp.tool = capture_tool
        register_server_tools(mcp)

        result = registered_tools["haproxy_remove_server"](
            domain="example.com",
            slot=99
        )

        assert "Error" in result
        assert "Slot" in result

    def test_remove_server_success(self, mock_socket_class, mock_select, patch_config_paths, response_builder):
        """Successfully remove server."""
        with open(patch_config_paths["map_file"], "w") as f:
            f.write("example.com pool_1\n")

        mock_sock = mock_socket_class(responses={
            "set server": "",
        })

        with patch("socket.socket", return_value=mock_sock):
            from haproxy_mcp.tools.servers import register_server_tools
            mcp = MagicMock()
            registered_tools = {}

            def capture_tool():
                def decorator(func):
                    registered_tools[func.__name__] = func
                    return func
                return decorator

            mcp.tool = capture_tool
            register_server_tools(mcp)

            result = registered_tools["haproxy_remove_server"](
                domain="example.com",
                slot=1
            )

        assert "Removed" in result
        assert "slot 1" in result


class TestHaproxySetServerState:
    """Tests for haproxy_set_server_state tool function."""

    def test_set_state_invalid_backend(self, patch_config_paths):
        """Reject invalid backend name."""
        from haproxy_mcp.tools.servers import register_server_tools
        mcp = MagicMock()
        registered_tools = {}

        def capture_tool():
            def decorator(func):
                registered_tools[func.__name__] = func
                return func
            return decorator

        mcp.tool = capture_tool
        register_server_tools(mcp)

        result = registered_tools["haproxy_set_server_state"](
            backend="invalid@backend",
            server="pool_1_1",
            state="ready"
        )

        assert "Error" in result
        assert "Invalid backend" in result

    def test_set_state_invalid_state(self, patch_config_paths):
        """Reject invalid state."""
        from haproxy_mcp.tools.servers import register_server_tools
        mcp = MagicMock()
        registered_tools = {}

        def capture_tool():
            def decorator(func):
                registered_tools[func.__name__] = func
                return func
            return decorator

        mcp.tool = capture_tool
        register_server_tools(mcp)

        result = registered_tools["haproxy_set_server_state"](
            backend="pool_1",
            server="pool_1_1",
            state="invalid"
        )

        assert "Error" in result
        assert "state must be" in result

    def test_set_state_success(self, mock_socket_class, mock_select, patch_config_paths):
        """Successfully set server state."""
        mock_sock = mock_socket_class(responses={
            "set server": "",
        })

        with patch("socket.socket", return_value=mock_sock):
            from haproxy_mcp.tools.servers import register_server_tools
            mcp = MagicMock()
            registered_tools = {}

            def capture_tool():
                def decorator(func):
                    registered_tools[func.__name__] = func
                    return func
                return decorator

            mcp.tool = capture_tool
            register_server_tools(mcp)

            result = registered_tools["haproxy_set_server_state"](
                backend="pool_1",
                server="pool_1_1",
                state="maint"
            )

        assert "pool_1/pool_1_1" in result
        assert "maint" in result


class TestHaproxySetServerWeight:
    """Tests for haproxy_set_server_weight tool function."""

    def test_set_weight_invalid_weight(self, patch_config_paths):
        """Reject invalid weight."""
        from haproxy_mcp.tools.servers import register_server_tools
        mcp = MagicMock()
        registered_tools = {}

        def capture_tool():
            def decorator(func):
                registered_tools[func.__name__] = func
                return func
            return decorator

        mcp.tool = capture_tool
        register_server_tools(mcp)

        result = registered_tools["haproxy_set_server_weight"](
            backend="pool_1",
            server="pool_1_1",
            weight=300  # > 256
        )

        assert "Error" in result
        assert "weight" in result

    def test_set_weight_success(self, mock_socket_class, mock_select, patch_config_paths):
        """Successfully set server weight."""
        mock_sock = mock_socket_class(responses={
            "set server": "",
        })

        with patch("socket.socket", return_value=mock_sock):
            from haproxy_mcp.tools.servers import register_server_tools
            mcp = MagicMock()
            registered_tools = {}

            def capture_tool():
                def decorator(func):
                    registered_tools[func.__name__] = func
                    return func
                return decorator

            mcp.tool = capture_tool
            register_server_tools(mcp)

            result = registered_tools["haproxy_set_server_weight"](
                backend="pool_1",
                server="pool_1_1",
                weight=2
            )

        assert "pool_1/pool_1_1" in result
        assert "2" in result


class TestConfigureServerSlot:
    """Tests for configure_server_slot helper function."""

    def test_configure_slot(self, mock_socket_class, mock_select):
        """Configure server slot sends correct commands."""
        mock_sock = mock_socket_class(responses={
            "set server": "",
        })

        with patch("socket.socket", return_value=mock_sock):
            from haproxy_mcp.tools.servers import configure_server_slot

            result = configure_server_slot("pool_1", "pool_1", 1, "10.0.0.1", 8080)

        assert result == "pool_1_1"
        # Verify commands were sent
        assert len(mock_sock.sent_commands) == 2
        assert "addr 10.0.0.1 port 8080" in mock_sock.sent_commands[0]
        assert "state ready" in mock_sock.sent_commands[1]


class TestHaproxyAddServersRollback:
    """Tests for haproxy_add_servers rollback functionality."""

    def test_add_servers_partial_failure_rollback(self, mock_socket_class, mock_select, patch_config_paths):
        """Rollback only failed slots when HAProxy error occurs."""
        with open(patch_config_paths["map_file"], "w") as f:
            f.write("example.com pool_1\n")

        # Mock configure_server_slot to fail on second slot
        call_count = [0]

        def mock_configure_server_slot(backend, server_prefix, slot, ip, http_port):
            call_count[0] += 1
            if slot == 2:
                raise HaproxyError("HAProxy command failed: server not found")
            return f"{server_prefix}_{slot}"

        mock_sock = mock_socket_class(responses={
            "set server": "",
        })

        with patch("socket.socket", return_value=mock_sock):
            with patch(
                "haproxy_mcp.tools.servers.configure_server_slot",
                side_effect=mock_configure_server_slot
            ):
                from haproxy_mcp.tools.servers import register_server_tools
                mcp = MagicMock()
                registered_tools = {}

                def capture_tool():
                    def decorator(func):
                        registered_tools[func.__name__] = func
                        return func
                    return decorator

                mcp.tool = capture_tool
                register_server_tools(mcp)

                result = registered_tools["haproxy_add_servers"](
                    domain="example.com",
                    servers='[{"slot": 1, "ip": "10.0.0.1"}, {"slot": 2, "ip": "10.0.0.2"}]'
                )

        # First server should be added, second should fail
        assert "Added 1 server" in result
        assert "Failed to add 1 server" in result
        assert "slot 1" in result  # Successfully added
        assert "slot 2" in result  # Failed

        # Verify servers.json only has successfully added server
        with open(patch_config_paths["servers_file"], "r") as f:
            config = json.load(f)
        assert "example.com" in config
        assert "1" in config["example.com"]  # Successfully added stays
        assert "2" not in config["example.com"]  # Failed one was rolled back

    def test_add_servers_unexpected_error_rollback_only_successful(
        self, mock_socket_class, mock_select, patch_config_paths
    ):
        """Rollback only successfully added servers on unexpected error."""
        with open(patch_config_paths["map_file"], "w") as f:
            f.write("example.com pool_1\n")

        # Track which servers were configured
        configured_slots = []

        # Mock socket that succeeds first then throws unexpected error
        original_configure = None

        def mock_configure_server_slot(backend, server_prefix, slot, ip, http_port):
            if slot == 2:
                # Simulate unexpected error (IOError is caught by the exception handler)
                raise IOError("Unexpected system error")
            configured_slots.append(slot)
            return f"{server_prefix}_{slot}"

        mock_sock = mock_socket_class(responses={
            "set server": "",
        })

        with patch("socket.socket", return_value=mock_sock):
            with patch(
                "haproxy_mcp.tools.servers.configure_server_slot",
                side_effect=mock_configure_server_slot
            ):
                from haproxy_mcp.tools.servers import register_server_tools
                mcp = MagicMock()
                registered_tools = {}

                def capture_tool():
                    def decorator(func):
                        registered_tools[func.__name__] = func
                        return func
                    return decorator

                mcp.tool = capture_tool
                register_server_tools(mcp)

                result = registered_tools["haproxy_add_servers"](
                    domain="example.com",
                    servers='[{"slot": 1, "ip": "10.0.0.1"}, {"slot": 2, "ip": "10.0.0.2"}, {"slot": 3, "ip": "10.0.0.3"}]'
                )

        # Should return error
        assert "Error" in result
        assert "Unexpected system error" in result

        # Verify servers.json is empty (all rolled back)
        with open(patch_config_paths["servers_file"], "r") as f:
            config = json.load(f)
        assert config == {} or "example.com" not in config or config.get("example.com") == {}

    def test_add_servers_rollback_failure_logged(
        self, mock_socket_class, mock_select, patch_config_paths, caplog
    ):
        """Log rollback failures during error recovery."""
        import logging
        with open(patch_config_paths["map_file"], "w") as f:
            f.write("example.com pool_1\n")

        def mock_configure_server_slot(backend, server_prefix, slot, ip, http_port):
            if slot == 2:
                raise OSError("Unexpected error")
            return f"{server_prefix}_{slot}"

        def mock_remove_server_from_config(domain, slot):
            raise IOError("Disk full")

        mock_sock = mock_socket_class(responses={
            "set server": "",
        })

        with patch("socket.socket", return_value=mock_sock):
            with patch(
                "haproxy_mcp.tools.servers.configure_server_slot",
                side_effect=mock_configure_server_slot
            ):
                with patch(
                    "haproxy_mcp.tools.servers.remove_server_from_config",
                    side_effect=mock_remove_server_from_config
                ):
                    with caplog.at_level(logging.ERROR, logger="haproxy_mcp"):
                        from haproxy_mcp.tools.servers import register_server_tools
                        mcp = MagicMock()
                        registered_tools = {}

                        def capture_tool():
                            def decorator(func):
                                registered_tools[func.__name__] = func
                                return func
                            return decorator

                        mcp.tool = capture_tool
                        register_server_tools(mcp)

                        result = registered_tools["haproxy_add_servers"](
                            domain="example.com",
                            servers='[{"slot": 1, "ip": "10.0.0.1"}, {"slot": 2, "ip": "10.0.0.2"}]'
                        )

        # Should return error from the original failure
        assert "Error" in result
        assert "Unexpected error" in result

        # Should have logged rollback failures
        assert any("Failed to rollback" in record.message for record in caplog.records)


class TestHaproxyAddServerAutoSlot:
    """Additional tests for auto-slot selection in haproxy_add_server."""

    def test_add_server_auto_slot_all_used(self, mock_socket_class, mock_select, patch_config_paths, response_builder):
        """Auto-select slot fails when all slots are in use."""
        with open(patch_config_paths["map_file"], "w") as f:
            f.write("example.com pool_1\n")

        # Build response with all 10 slots used
        servers = []
        for i in range(1, 11):  # MAX_SLOTS = 10
            servers.append({
                "be_name": "pool_1",
                "srv_name": f"pool_1_{i}",
                "srv_addr": f"10.0.0.{i}",
                "srv_port": 80
            })

        mock_sock = mock_socket_class(responses={
            "show servers state": response_builder.servers_state(servers),
            "set server": "",
        })

        with patch("socket.socket", return_value=mock_sock):
            from haproxy_mcp.tools.servers import register_server_tools
            mcp = MagicMock()
            registered_tools = {}

            def capture_tool():
                def decorator(func):
                    registered_tools[func.__name__] = func
                    return func
                return decorator

            mcp.tool = capture_tool
            register_server_tools(mcp)

            result = registered_tools["haproxy_add_server"](
                domain="example.com",
                slot=0,  # Auto-select
                ip="10.0.0.100",
                http_port=80
            )

        assert "Error" in result
        assert "No available slots" in result

    def test_add_server_negative_slot_auto_select(self, mock_socket_class, mock_select, patch_config_paths, response_builder):
        """Negative slot number triggers auto-selection."""
        with open(patch_config_paths["map_file"], "w") as f:
            f.write("example.com pool_1\n")

        mock_sock = mock_socket_class(responses={
            "show servers state": response_builder.servers_state([
                {"be_name": "pool_1", "srv_name": "pool_1_1", "srv_addr": "0.0.0.0", "srv_port": 0},
            ]),
            "set server": "",
        })

        with patch("socket.socket", return_value=mock_sock):
            from haproxy_mcp.tools.servers import register_server_tools
            mcp = MagicMock()
            registered_tools = {}

            def capture_tool():
                def decorator(func):
                    registered_tools[func.__name__] = func
                    return func
                return decorator

            mcp.tool = capture_tool
            register_server_tools(mcp)

            result = registered_tools["haproxy_add_server"](
                domain="example.com",
                slot=-1,  # Negative triggers auto-select
                ip="10.0.0.1",
                http_port=80
            )

        assert "slot 1" in result  # First available slot


class TestHaproxyAddServersPartialFailure:
    """Tests for partial failure scenarios in haproxy_add_servers."""

    def test_add_servers_validation_error_per_server(self, patch_config_paths):
        """Handle validation errors for individual servers."""
        from haproxy_mcp.tools.servers import register_server_tools
        mcp = MagicMock()
        registered_tools = {}

        def capture_tool():
            def decorator(func):
                registered_tools[func.__name__] = func
                return func
            return decorator

        mcp.tool = capture_tool
        register_server_tools(mcp)

        # Mix of valid and invalid servers
        result = registered_tools["haproxy_add_servers"](
            domain="example.com",
            servers='[{"slot": 1, "ip": "10.0.0.1"}, {"slot": "invalid"}, {"slot": 2}]'
        )

        assert "Validation errors" in result
        assert "Server 2" in result  # Invalid slot type
        assert "Server 3" in result  # Missing IP


class TestHaproxyWaitDrain:
    """Tests for haproxy_wait_drain tool function."""

    def test_wait_drain_success(self, patch_config_paths):
        """Successfully wait for connections to drain."""
        with open(patch_config_paths["map_file"], "w") as f:
            f.write("example.com pool_1\n")

        # Mock haproxy_cmd to return 0 connections
        with patch("haproxy_mcp.tools.servers.haproxy_cmd") as mock_cmd:
            mock_cmd.return_value = "# pxname,svname,qcur,qmax,scur,smax\npool_1,pool_1_1,0,0,0,0\n"
            # Provide enough time values: start_time, loop check, elapsed calculation
            with patch("time.time", side_effect=[0, 0.1, 0.2]):
                from haproxy_mcp.tools.servers import register_server_tools
                mcp = MagicMock()
                registered_tools = {}

                def capture_tool():
                    def decorator(func):
                        registered_tools[func.__name__] = func
                        return func
                    return decorator

                mcp.tool = capture_tool
                register_server_tools(mcp)

                result = registered_tools["haproxy_wait_drain"](
                    domain="example.com",
                    timeout=30
                )

        assert "drained" in result.lower()

    def test_wait_drain_timeout(self, patch_config_paths):
        """Timeout when connections don't drain."""
        with open(patch_config_paths["map_file"], "w") as f:
            f.write("example.com pool_1\n")

        time_values = [0, 0.5, 1.0, 1.5, 2.0, 2.5, 3.0]  # Simulate time passing
        time_iter = iter(time_values)

        # Mock haproxy_cmd to return active connections
        with patch("haproxy_mcp.tools.servers.haproxy_cmd") as mock_cmd:
            mock_cmd.return_value = "# pxname,svname,qcur,qmax,scur,smax\npool_1,pool_1_1,0,0,5,10\n"
            with patch("time.time", side_effect=lambda: next(time_iter)):
                with patch("time.sleep", return_value=None):  # Don't actually sleep
                    from haproxy_mcp.tools.servers import register_server_tools
                    mcp = MagicMock()
                    registered_tools = {}

                    def capture_tool():
                        def decorator(func):
                            registered_tools[func.__name__] = func
                            return func
                        return decorator

                    mcp.tool = capture_tool
                    register_server_tools(mcp)

                    result = registered_tools["haproxy_wait_drain"](
                        domain="example.com",
                        timeout=2  # Short timeout
                    )

        assert "Timeout" in result

    def test_wait_drain_invalid_domain(self, patch_config_paths):
        """Reject invalid domain format."""
        from haproxy_mcp.tools.servers import register_server_tools
        mcp = MagicMock()
        registered_tools = {}

        def capture_tool():
            def decorator(func):
                registered_tools[func.__name__] = func
                return func
            return decorator

        mcp.tool = capture_tool
        register_server_tools(mcp)

        result = registered_tools["haproxy_wait_drain"](
            domain="-invalid",
            timeout=30
        )

        assert "Error" in result
        assert "Invalid domain" in result

    def test_wait_drain_invalid_timeout(self, patch_config_paths):
        """Reject invalid timeout values."""
        from haproxy_mcp.tools.servers import register_server_tools
        mcp = MagicMock()
        registered_tools = {}

        def capture_tool():
            def decorator(func):
                registered_tools[func.__name__] = func
                return func
            return decorator

        mcp.tool = capture_tool
        register_server_tools(mcp)

        # Test timeout > 300
        result = registered_tools["haproxy_wait_drain"](
            domain="example.com",
            timeout=500
        )
        assert "Error" in result
        assert "Timeout must be" in result

        # Test timeout < 1
        result = registered_tools["haproxy_wait_drain"](
            domain="example.com",
            timeout=0
        )
        assert "Error" in result

    def test_wait_drain_domain_not_found(self, mock_socket_class, mock_select, patch_config_paths):
        """Error when domain not found in map."""
        # Empty map file - domain not configured
        with open(patch_config_paths["map_file"], "w") as f:
            f.write("")

        from haproxy_mcp.tools.servers import register_server_tools
        mcp = MagicMock()
        registered_tools = {}

        def capture_tool():
            def decorator(func):
                registered_tools[func.__name__] = func
                return func
            return decorator

        mcp.tool = capture_tool
        register_server_tools(mcp)

        result = registered_tools["haproxy_wait_drain"](
            domain="unknown.com",
            timeout=30
        )

        assert "Error" in result


class TestHaproxySetServerWeightBoundary:
    """Tests for haproxy_set_server_weight boundary values."""

    def test_set_weight_zero(self, mock_socket_class, mock_select, patch_config_paths):
        """Set server weight to 0 (disabled)."""
        mock_sock = mock_socket_class(responses={
            "set server": "",
        })

        with patch("socket.socket", return_value=mock_sock):
            from haproxy_mcp.tools.servers import register_server_tools
            mcp = MagicMock()
            registered_tools = {}

            def capture_tool():
                def decorator(func):
                    registered_tools[func.__name__] = func
                    return func
                return decorator

            mcp.tool = capture_tool
            register_server_tools(mcp)

            result = registered_tools["haproxy_set_server_weight"](
                backend="pool_1",
                server="pool_1_1",
                weight=0
            )

        assert "pool_1/pool_1_1" in result
        assert "0" in result

    def test_set_weight_max(self, mock_socket_class, mock_select, patch_config_paths):
        """Set server weight to maximum 256."""
        mock_sock = mock_socket_class(responses={
            "set server": "",
        })

        with patch("socket.socket", return_value=mock_sock):
            from haproxy_mcp.tools.servers import register_server_tools
            mcp = MagicMock()
            registered_tools = {}

            def capture_tool():
                def decorator(func):
                    registered_tools[func.__name__] = func
                    return func
                return decorator

            mcp.tool = capture_tool
            register_server_tools(mcp)

            result = registered_tools["haproxy_set_server_weight"](
                backend="pool_1",
                server="pool_1_1",
                weight=256
            )

        assert "pool_1/pool_1_1" in result
        assert "256" in result

    def test_set_weight_negative(self, patch_config_paths):
        """Reject negative weight."""
        from haproxy_mcp.tools.servers import register_server_tools
        mcp = MagicMock()
        registered_tools = {}

        def capture_tool():
            def decorator(func):
                registered_tools[func.__name__] = func
                return func
            return decorator

        mcp.tool = capture_tool
        register_server_tools(mcp)

        result = registered_tools["haproxy_set_server_weight"](
            backend="pool_1",
            server="pool_1_1",
            weight=-1
        )

        assert "Error" in result
        assert "weight" in result


class TestHaproxySetDomainState:
    """Tests for haproxy_set_domain_state tool function."""

    def test_set_domain_state_success(self, mock_socket_class, mock_select, patch_config_paths, response_builder):
        """Set all servers of a domain to a state."""
        with open(patch_config_paths["map_file"], "w") as f:
            f.write("example.com pool_1\n")

        mock_sock = mock_socket_class(responses={
            "show servers state": response_builder.servers_state([
                {"be_name": "pool_1", "srv_name": "pool_1_1", "srv_addr": "10.0.0.1", "srv_port": 80},
                {"be_name": "pool_1", "srv_name": "pool_1_2", "srv_addr": "10.0.0.2", "srv_port": 80},
            ]),
            "set server": "",
        })

        with patch("socket.socket", return_value=mock_sock):
            from haproxy_mcp.tools.servers import register_server_tools
            mcp = MagicMock()
            registered_tools = {}

            def capture_tool():
                def decorator(func):
                    registered_tools[func.__name__] = func
                    return func
                return decorator

            mcp.tool = capture_tool
            register_server_tools(mcp)

            result = registered_tools["haproxy_set_domain_state"](
                domain="example.com",
                state="maint"
            )

        assert "Set 2 servers" in result
        assert "maint" in result

    def test_set_domain_state_invalid_domain(self, patch_config_paths):
        """Reject invalid domain format."""
        from haproxy_mcp.tools.servers import register_server_tools
        mcp = MagicMock()
        registered_tools = {}

        def capture_tool():
            def decorator(func):
                registered_tools[func.__name__] = func
                return func
            return decorator

        mcp.tool = capture_tool
        register_server_tools(mcp)

        result = registered_tools["haproxy_set_domain_state"](
            domain="-invalid",
            state="ready"
        )

        assert "Error" in result
        assert "Invalid domain" in result

    def test_set_domain_state_invalid_state(self, patch_config_paths):
        """Reject invalid state value."""
        from haproxy_mcp.tools.servers import register_server_tools
        mcp = MagicMock()
        registered_tools = {}

        def capture_tool():
            def decorator(func):
                registered_tools[func.__name__] = func
                return func
            return decorator

        mcp.tool = capture_tool
        register_server_tools(mcp)

        result = registered_tools["haproxy_set_domain_state"](
            domain="example.com",
            state="invalid"
        )

        assert "Error" in result
        assert "must be" in result

    def test_set_domain_state_no_active_servers(self, mock_socket_class, mock_select, patch_config_paths, response_builder):
        """No active servers found for domain."""
        with open(patch_config_paths["map_file"], "w") as f:
            f.write("example.com pool_1\n")

        # All servers have 0.0.0.0 address (not configured)
        mock_sock = mock_socket_class(responses={
            "show servers state": response_builder.servers_state([
                {"be_name": "pool_1", "srv_name": "pool_1_1", "srv_addr": "0.0.0.0", "srv_port": 0},
            ]),
            "set server": "",
        })

        with patch("socket.socket", return_value=mock_sock):
            from haproxy_mcp.tools.servers import register_server_tools
            mcp = MagicMock()
            registered_tools = {}

            def capture_tool():
                def decorator(func):
                    registered_tools[func.__name__] = func
                    return func
                return decorator

            mcp.tool = capture_tool
            register_server_tools(mcp)

            result = registered_tools["haproxy_set_domain_state"](
                domain="example.com",
                state="ready"
            )

        assert "No active servers found" in result

    def test_set_domain_state_domain_not_found(self, mock_socket_class, mock_select, patch_config_paths, response_builder):
        """Handle domain not found in map - shows no active servers."""
        # Empty map file
        with open(patch_config_paths["map_file"], "w") as f:
            f.write("")

        # Mock should show no servers for unknown domain's backend
        mock_sock = mock_socket_class(responses={
            "show servers state": response_builder.servers_state([]),
        })

        with patch("socket.socket", return_value=mock_sock):
            from haproxy_mcp.tools.servers import register_server_tools
            mcp = MagicMock()
            registered_tools = {}

            def capture_tool():
                def decorator(func):
                    registered_tools[func.__name__] = func
                    return func
                return decorator

            mcp.tool = capture_tool
            register_server_tools(mcp)

            result = registered_tools["haproxy_set_domain_state"](
                domain="unknown.com",
                state="ready"
            )

        # When domain is not in map, get_backend_and_prefix raises ValueError
        # which is caught and returns Error
        assert "Error" in result or "No active servers" in result