test_server.py•9.65 kB
"""Test suite for NetBox Infrastructure MCP Server."""
import pytest
from unittest.mock import Mock, patch, MagicMock
from src.vault_client import VaultClient
from src.netbox_client import NetBoxClient
from src.state_confidence import StateConfidenceClient
from src.tools import hosts, virtual_machines, ip_addresses, vlans
@pytest.fixture
def mock_vault_client():
"""Mock Vault client."""
client = Mock(spec=VaultClient)
client._token = "test-token"
client._token_expiry = None
client.authenticate.return_value = True
client.mint_netbox_token.return_value = "netbox-api-token"
client.is_authenticated.return_value = True
return client
@pytest.fixture
def mock_netbox_client(mock_vault_client):
"""Mock NetBox client."""
client = Mock(spec=NetBoxClient)
client.vault_client = mock_vault_client
client.list_devices.return_value = [
{
"id": 1,
"name": "test-host",
"display": "test-host",
"primary_ip4": {"address": "192.168.1.10/24"},
"device_role": {"slug": "server"},
}
]
client.get_device.return_value = {
"id": 1,
"name": "test-host",
"display": "test-host",
}
client.search_devices.return_value = [
{
"id": 1,
"name": "test-host",
"display": "test-host",
}
]
client.list_virtual_machines.return_value = [
{
"id": 1,
"name": "test-vm",
"display": "test-vm",
"primary_ip4": {"address": "192.168.1.20/24"},
}
]
client.get_virtual_machine.return_value = {
"id": 1,
"name": "test-vm",
"display": "test-vm",
}
client.list_vm_interfaces.return_value = [
{
"id": 1,
"name": "eth0",
"ip_addresses": [
{"address": "192.168.1.20/24", "id": 1}
],
}
]
client.list_ip_addresses.return_value = [
{
"id": 1,
"address": "192.168.1.10/24",
"display": "192.168.1.10/24",
}
]
client.get_ip_address.return_value = {
"id": 1,
"address": "192.168.1.10/24",
"display": "192.168.1.10/24",
}
client.search_ip_addresses.return_value = [
{
"id": 1,
"address": "192.168.1.10/24",
"display": "192.168.1.10/24",
}
]
client.list_vlans.return_value = [
{
"id": 1,
"vid": 100,
"name": "test-vlan",
"display": "VLAN 100",
}
]
client.get_vlan.return_value = {
"id": 1,
"vid": 100,
"name": "test-vlan",
"display": "VLAN 100",
}
client.list_vlan_ip_addresses.return_value = [
{
"id": 1,
"address": "192.168.1.10/24",
"display": "192.168.1.10/24",
}
]
return client
@pytest.fixture
def mock_state_client():
"""Mock state confidence client."""
client = Mock(spec=StateConfidenceClient)
client.get_certainty_score.return_value = {
"certainty_score": 0.95,
"data_age_seconds": 300,
"last_verified": "2025-01-01T00:00:00Z",
}
return client
@pytest.mark.asyncio
async def test_list_hosts(mock_netbox_client, mock_state_client):
"""Test list_hosts tool."""
arguments = {"limit": 10, "include_certainty": True}
result = await hosts.handle_list_hosts(
arguments, mock_netbox_client, mock_state_client
)
assert len(result) == 1
assert "Found 1 host(s)" in result[0].text
mock_netbox_client.list_devices.assert_called_once()
@pytest.mark.asyncio
async def test_get_host(mock_netbox_client, mock_state_client):
"""Test get_host tool."""
arguments = {"hostname": "test-host", "include_certainty": True}
result = await hosts.handle_get_host(
arguments, mock_netbox_client, mock_state_client
)
assert len(result) == 1
assert "test-host" in result[0].text
mock_netbox_client.get_device.assert_called_once_with("test-host")
@pytest.mark.asyncio
async def test_search_hosts(mock_netbox_client, mock_state_client):
"""Test search_hosts tool."""
arguments = {"query": "test", "limit": 10, "include_certainty": True}
result = await hosts.handle_search_hosts(
arguments, mock_netbox_client, mock_state_client
)
assert len(result) == 1
assert "Found 1 matching host(s)" in result[0].text
mock_netbox_client.search_devices.assert_called_once()
@pytest.mark.asyncio
async def test_list_vms(mock_netbox_client, mock_state_client):
"""Test list_vms tool."""
arguments = {"limit": 10, "include_certainty": True}
result = await virtual_machines.handle_list_vms(
arguments, mock_netbox_client, mock_state_client
)
assert len(result) == 1
assert "Found 1 VM/container(s)" in result[0].text
mock_netbox_client.list_virtual_machines.assert_called_once()
@pytest.mark.asyncio
async def test_get_vm(mock_netbox_client, mock_state_client):
"""Test get_vm tool."""
arguments = {"hostname": "test-vm", "include_certainty": True}
result = await virtual_machines.handle_get_vm(
arguments, mock_netbox_client, mock_state_client
)
assert len(result) == 1
assert "test-vm" in result[0].text
mock_netbox_client.get_virtual_machine.assert_called_once_with("test-vm")
@pytest.mark.asyncio
async def test_list_vm_interfaces(mock_netbox_client, mock_state_client):
"""Test list_vm_interfaces tool."""
arguments = {"vm_name": "test-vm", "include_certainty": True}
result = await virtual_machines.handle_list_vm_interfaces(
arguments, mock_netbox_client, mock_state_client
)
assert len(result) == 1
assert "interface(s)" in result[0].text
mock_netbox_client.list_vm_interfaces.assert_called_once_with("test-vm")
@pytest.mark.asyncio
async def test_list_ips(mock_netbox_client, mock_state_client):
"""Test list_ips tool."""
arguments = {"limit": 10, "include_certainty": True}
result = await ip_addresses.handle_list_ips(
arguments, mock_netbox_client, mock_state_client
)
assert len(result) == 1
assert "Found 1 IP address(es)" in result[0].text
mock_netbox_client.list_ip_addresses.assert_called_once()
@pytest.mark.asyncio
async def test_get_ip(mock_netbox_client, mock_state_client):
"""Test get_ip tool."""
arguments = {"ip_address": "192.168.1.10/24", "include_certainty": True}
result = await ip_addresses.handle_get_ip(
arguments, mock_netbox_client, mock_state_client
)
assert len(result) == 1
assert "192.168.1.10" in result[0].text
mock_netbox_client.get_ip_address.assert_called_once()
@pytest.mark.asyncio
async def test_search_ips(mock_netbox_client, mock_state_client):
"""Test search_ips tool."""
arguments = {"query": "192.168.1", "limit": 10, "include_certainty": True}
result = await ip_addresses.handle_search_ips(
arguments, mock_netbox_client, mock_state_client
)
assert len(result) == 1
assert "Found 1 matching IP address(es)" in result[0].text
mock_netbox_client.search_ip_addresses.assert_called_once()
@pytest.mark.asyncio
async def test_list_vlans(mock_netbox_client, mock_state_client):
"""Test list_vlans tool."""
arguments = {"limit": 10, "include_certainty": True}
result = await vlans.handle_list_vlans(
arguments, mock_netbox_client, mock_state_client
)
assert len(result) == 1
assert "Found 1 VLAN(s)" in result[0].text
mock_netbox_client.list_vlans.assert_called_once()
@pytest.mark.asyncio
async def test_get_vlan(mock_netbox_client, mock_state_client):
"""Test get_vlan tool."""
arguments = {"vlan_id": 100, "include_certainty": True}
result = await vlans.handle_get_vlan(
arguments, mock_netbox_client, mock_state_client
)
assert len(result) == 1
assert "VLAN 100" in result[0].text or "100" in result[0].text
mock_netbox_client.get_vlan.assert_called_once_with(100)
@pytest.mark.asyncio
async def test_list_vlan_ips(mock_netbox_client, mock_state_client):
"""Test list_vlan_ips tool."""
arguments = {"vlan_id": 100, "include_certainty": True}
result = await vlans.handle_list_vlan_ips(
arguments, mock_netbox_client, mock_state_client
)
assert len(result) == 1
assert "IP address(es) in VLAN 100" in result[0].text
mock_netbox_client.list_vlan_ip_addresses.assert_called_once_with(100)
@pytest.mark.asyncio
async def test_missing_hostname():
"""Test error handling for missing hostname."""
mock_client = Mock()
arguments = {}
result = await hosts.handle_get_host(arguments, mock_client, None)
assert len(result) == 1
assert "Error" in result[0].text or "required" in result[0].text
@pytest.mark.asyncio
async def test_host_not_found(mock_netbox_client):
"""Test handling when host is not found."""
mock_netbox_client.get_device.return_value = None
arguments = {"hostname": "nonexistent"}
result = await hosts.handle_get_host(arguments, mock_netbox_client, None)
assert len(result) == 1
assert "not found" in result[0].text.lower()
def test_vault_client_authentication(mock_vault_client):
"""Test Vault client authentication."""
assert mock_vault_client.authenticate() is True
def test_vault_client_mint_token(mock_vault_client):
"""Test Vault client token minting."""
token = mock_vault_client.mint_netbox_token("netbox/jit-tokens")
assert token == "netbox-api-token"