"""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