"""Streamlined core functionality tests for 1.0 release.
This module contains only the essential tests that prove core functionality
works without external dependencies or complex edge cases.
"""
import asyncio
from unittest.mock import Mock
import pytest
from mcp.types import TextContent
from src.server import list_tools
from src.structured_protocol import ProtocolOptimizer, StructuredProtocol
from src.tiny_llm import TinyLLM, TinyLLMClient
class TestCoreMCPServer:
"""Core MCP server functionality tests."""
@pytest.mark.asyncio
async def test_tool_discovery(self):
"""Test that tools can be discovered."""
tools = await list_tools()
assert len(tools) > 0
assert all(hasattr(tool, "name") for tool in tools)
assert all(hasattr(tool, "description") for tool in tools)
@pytest.mark.asyncio
async def test_tool_categories(self):
"""Test that all expected tool categories are present."""
tools = await list_tools()
tool_names = [tool.name for tool in tools]
expected_categories = {
"hosts": ["list_hosts", "get_host", "search_hosts"],
"vms": ["list_vms", "get_vm", "list_vm_interfaces"],
"ips": ["list_ips", "get_ip", "search_ips"],
"vlans": ["list_vlans", "get_vlan", "list_vlan_ips"],
}
for category, expected_tools in expected_categories.items():
for tool_name in expected_tools:
assert (
tool_name in tool_names
), f"Missing {category} tool: {tool_name}"
@pytest.mark.asyncio
async def test_tool_execution_with_mocks(self):
"""Test tool execution with mocked dependencies."""
# Mock the clients
mock_netbox_client = Mock()
mock_netbox_client.list_devices.return_value = [
{"id": 1, "name": "test-server", "status": "active"}
]
mock_state_client = Mock()
mock_state_client.get_certainty_score.return_value = 0.95
# Test list_hosts tool
from src.tools.hosts import handle_list_hosts
result = await handle_list_hosts(
{"limit": 5, "include_certainty": True},
mock_netbox_client,
mock_state_client,
)
assert len(result) > 0
assert isinstance(result[0], TextContent)
assert "test-server" in result[0].text
class TestTinyLLMCore:
"""Core Tiny LLM functionality tests."""
@pytest.mark.asyncio
async def test_basic_response_generation(self):
"""Test basic response generation."""
llm = TinyLLM("test-llm", base_delay=0.01, max_delay=0.02)
response = await llm.generate_response(
"Test query", tools=None, max_tokens=50
)
assert response.content is not None
assert response.tokens_used > 0
assert response.processing_time > 0
assert 0.0 <= response.confidence <= 1.0
@pytest.mark.asyncio
async def test_tool_call_generation(self):
"""Test tool call generation."""
llm = TinyLLM("test-llm", base_delay=0.01, max_delay=0.02)
tools = [{"name": "list_hosts", "category": "hosts"}]
tool_calls = await llm.generate_tool_calls(
"Show me all servers", tools
)
assert isinstance(tool_calls, list)
assert len(tool_calls) >= 0
@pytest.mark.asyncio
async def test_client_api(self):
"""Test client API functionality."""
client = TinyLLMClient("test-client")
messages = [{"role": "user", "content": "Test query"}]
response = await client.chat_completion(messages, max_tokens=50)
assert "choices" in response
assert "usage" in response
assert response["choices"][0]["message"]["role"] == "assistant"
@pytest.mark.asyncio
async def test_concurrent_requests(self):
"""Test concurrent request handling."""
client = TinyLLMClient("test-client")
tasks = []
for i in range(3):
messages = [{"role": "user", "content": f"Test query {i}"}]
task = client.chat_completion(messages, max_tokens=50)
tasks.append(task)
responses = await asyncio.gather(*tasks)
assert len(responses) == 3
assert all("choices" in r for r in responses)
class TestStructuredProtocol:
"""Structured protocol functionality tests."""
def test_query_creation(self):
"""Test structured query creation."""
protocol = StructuredProtocol()
query = protocol.create_query(
user_id="test_user",
query="Show me all servers",
context={"test": True},
)
assert query.user_id == "test_user"
assert query.query == "Show me all servers"
assert query.context["test"] is True
def test_router_decision(self):
"""Test router decision optimization."""
protocol = StructuredProtocol()
optimizer = ProtocolOptimizer()
query = protocol.create_query(
user_id="test_user", query="Show me all servers"
)
decision = optimizer.optimize_router_decision(query)
assert decision.model_location is not None
assert isinstance(decision.tools_needed, list)
assert decision.priority >= 1
def test_llm_request_optimization(self):
"""Test LLM request optimization."""
protocol = StructuredProtocol()
optimizer = ProtocolOptimizer()
query = protocol.create_query(
user_id="test_user", query="Show me all servers"
)
decision = optimizer.optimize_router_decision(query)
llm_request = optimizer.optimize_llm_request(query, decision)
assert isinstance(llm_request.tools_available, list)
assert llm_request.max_tokens > 0
assert llm_request.temperature >= 0.0
class TestEndToEndWorkflow:
"""End-to-end workflow tests."""
@pytest.mark.asyncio
async def test_complete_workflow(self):
"""Test complete workflow from query to response."""
protocol = StructuredProtocol()
optimizer = ProtocolOptimizer()
client = TinyLLMClient("test-client")
# Create query
query = protocol.create_query(
user_id="test_user", query="Show me all web servers"
)
# Router decision
decision = optimizer.optimize_router_decision(query)
# LLM request
llm_request = optimizer.optimize_llm_request(query, decision)
# LLM processing
messages = [{"role": "user", "content": query.query}]
llm_response = await client.chat_completion(
messages=messages,
tools=llm_request.tools_available,
max_tokens=llm_request.max_tokens,
)
# Verify response
assert llm_response["choices"][0]["message"]["role"] == "assistant"
assert llm_response["usage"]["total_tokens"] > 0
@pytest.mark.asyncio
async def test_error_handling(self):
"""Test error handling in workflows."""
client = TinyLLMClient("test-client")
# Test with empty messages
try:
await client.chat_completion([], max_tokens=50)
assert False, "Should have raised ValueError"
except ValueError:
pass # Expected
@pytest.mark.asyncio
async def test_edge_cases(self):
"""Test edge cases."""
llm = TinyLLM("test-llm", base_delay=0.01, max_delay=0.02)
# Test empty prompt
response = await llm.generate_response("", tools=None, max_tokens=50)
assert response.content is not None
assert response.tokens_used >= 0
# Test long prompt
long_prompt = "x" * 1000
response = await llm.generate_response(
long_prompt, tools=None, max_tokens=100
)
assert response.content is not None
assert response.tokens_used > 0
class TestPerformanceBasics:
"""Basic performance tests."""
@pytest.mark.asyncio
async def test_response_time(self):
"""Test response time is reasonable."""
llm = TinyLLM("perf-test", base_delay=0.1, max_delay=0.2)
start_time = asyncio.get_event_loop().time()
response = await llm.generate_response(
"Performance test", tools=None, max_tokens=50
)
end_time = asyncio.get_event_loop().time()
total_time = end_time - start_time
assert total_time < 1.0, f"Response too slow: {total_time:.3f}s"
assert response.processing_time > 0
@pytest.mark.asyncio
async def test_memory_efficiency(self):
"""Test memory efficiency."""
client = TinyLLMClient("test-client")
# Process multiple requests
for i in range(10):
messages = [{"role": "user", "content": f"Test query {i}"}]
response = await client.chat_completion(messages, max_tokens=50)
assert "choices" in response
# If we get here without memory issues, test passes
assert True
