Files
resolutionflow/backend/tests/test_ai_endpoints.py
chihlasm 44432413c2 feat: AI-assisted flow builder with 4-stage wizard
Implements the complete AI flow builder feature using a guided 4-stage
wizard (Foundation → Scaffold → Branch Detail → Review & Assemble).
AI assists at bounded points using Claude Haiku for cost-efficient
structured JSON generation (~$0.01-0.03/flow).

Backend: new models (ai_conversations, ai_usage), Alembic migration,
quota enforcement with billing anchor, Anthropic API integration with
prompt caching, tree validation, conversation CRUD with 24h TTL,
APScheduler cleanup job, 5 API endpoints, Pydantic schemas.

Frontend: TypeScript types, API client, Zustand store for wizard state,
7 components (modal, step indicator, foundation form, branch selector,
branch detail view, tree preview, quota display), MyTreesPage integration
with "Build with AI" button (hidden when AI not configured).

Tests: 14 validator unit tests + 11 endpoint integration tests with
mocked Anthropic (zero real API spend). All 25 tests passing.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
2026-02-20 08:07:08 -05:00

361 lines
12 KiB
Python

"""Integration tests for AI Flow Builder endpoints.
All Anthropic API calls are mocked — zero real API spend.
"""
import json
from unittest.mock import AsyncMock, patch, MagicMock
import pytest
from app.core.config import settings
# ── Sample AI responses ──
SCAFFOLD_RESPONSE_JSON = json.dumps({
"branches": [
{"name": "Service Not Running", "description": "The target service is stopped or crashed."},
{"name": "Authentication Failures", "description": "Users cannot authenticate against the service."},
{"name": "Network Connectivity", "description": "Network-level issues preventing access."},
{"name": "Configuration Errors", "description": "Misconfiguration of the service or its dependencies."},
]
})
BRANCH_DETAIL_JSON = json.dumps({
"id": "svc-root",
"type": "decision",
"question": "Is the service running?",
"options": [
{"id": "opt-yes", "label": "Yes", "next_node_id": "svc-check-logs"},
{"id": "opt-no", "label": "No", "next_node_id": "svc-restart"},
],
"children": [
{
"id": "svc-check-logs",
"type": "action",
"title": "Check Event Logs",
"description": "Check Windows Event Viewer for errors.",
"commands": ["Get-EventLog -LogName Application -Newest 20"],
"children": [
{
"id": "svc-logs-resolved",
"type": "solution",
"title": "Issue Found in Logs",
"description": "Error identified and resolved.",
"resolution_steps": ["Fix the error", "Restart service"],
}
],
},
{
"id": "svc-restart",
"type": "action",
"title": "Restart Service",
"description": "Attempt to restart the service.",
"commands": ["Restart-Service -Name 'TestService'"],
"children": [
{
"id": "svc-restart-ok",
"type": "solution",
"title": "Service Restored",
"description": "Service is running after restart.",
"resolution_steps": ["Verify connectivity", "Document in ticket"],
}
],
},
],
})
def _mock_anthropic_response(text: str, input_tokens: int = 100, output_tokens: int = 200):
"""Create a mock Anthropic API response."""
response = MagicMock()
response.content = [MagicMock(text=text)]
response.usage = MagicMock(input_tokens=input_tokens, output_tokens=output_tokens)
return response
@pytest.fixture
def enable_ai():
"""Temporarily enable AI by setting a fake API key."""
original = settings.ANTHROPIC_API_KEY
settings.ANTHROPIC_API_KEY = "test-key-fake"
yield
settings.ANTHROPIC_API_KEY = original
@pytest.fixture
def disable_ai():
"""Ensure AI is disabled."""
original = settings.ANTHROPIC_API_KEY
settings.ANTHROPIC_API_KEY = None
yield
settings.ANTHROPIC_API_KEY = original
# ── Quota endpoint ──
@pytest.mark.asyncio
async def test_quota_returns_disabled_when_no_key(client, auth_headers, disable_ai):
"""GET /ai/quota returns ai_enabled=false when no API key."""
response = await client.get("/api/v1/ai/quota", headers=auth_headers)
assert response.status_code == 200
data = response.json()
assert data["ai_enabled"] is False
assert data["allowed"] is False
@pytest.mark.asyncio
async def test_quota_returns_enabled_with_key(client, auth_headers, enable_ai):
"""GET /ai/quota returns ai_enabled=true with API key configured."""
response = await client.get("/api/v1/ai/quota", headers=auth_headers)
assert response.status_code == 200
data = response.json()
assert data["ai_enabled"] is True
assert data["allowed"] is True
# ── Start endpoint ──
@pytest.mark.asyncio
async def test_start_requires_auth(client, enable_ai):
"""POST /ai/start requires authentication."""
response = await client.post("/api/v1/ai/start", json={
"flow_type": "troubleshooting",
"name": "Test Flow",
"description": "Test",
})
assert response.status_code == 401
@pytest.mark.asyncio
async def test_start_returns_503_when_disabled(client, auth_headers, disable_ai):
"""POST /ai/start returns 503 when AI is not configured."""
response = await client.post(
"/api/v1/ai/start",
json={
"flow_type": "troubleshooting",
"name": "Test Flow",
"description": "Test description",
},
headers=auth_headers,
)
assert response.status_code == 503
@pytest.mark.asyncio
async def test_start_creates_conversation(client, auth_headers, enable_ai):
"""POST /ai/start creates a conversation and returns conversation_id."""
response = await client.post(
"/api/v1/ai/start",
json={
"flow_type": "troubleshooting",
"name": "DNS Issues",
"description": "Troubleshooting DNS resolution failures",
"environment_tags": ["Windows Server", "Active Directory"],
},
headers=auth_headers,
)
assert response.status_code == 201
data = response.json()
assert "conversation_id" in data
assert data["status"] == "foundation"
@pytest.mark.asyncio
async def test_start_validates_input(client, auth_headers, enable_ai):
"""POST /ai/start rejects invalid input."""
response = await client.post(
"/api/v1/ai/start",
json={
"flow_type": "troubleshooting",
"name": "", # Empty name
"description": "Test",
},
headers=auth_headers,
)
assert response.status_code == 422
# ── Scaffold endpoint ──
@pytest.mark.asyncio
async def test_scaffold_success(client, auth_headers, enable_ai):
"""POST /ai/scaffold returns AI-generated branches."""
# Create conversation first
start_resp = await client.post(
"/api/v1/ai/start",
json={
"flow_type": "troubleshooting",
"name": "DNS Issues",
"description": "DNS resolution failures",
},
headers=auth_headers,
)
conversation_id = start_resp.json()["conversation_id"]
# Mock Anthropic
mock_response = _mock_anthropic_response(SCAFFOLD_RESPONSE_JSON)
with patch("app.core.ai_tree_generator_service._get_client") as mock_client:
mock_client.return_value.messages.create = AsyncMock(return_value=mock_response)
response = await client.post(
"/api/v1/ai/scaffold",
json={"conversation_id": conversation_id},
headers=auth_headers,
)
assert response.status_code == 200
data = response.json()
assert data["status"] == "scaffolding"
assert len(data["branches"]) == 4
assert data["branches"][0]["name"] == "Service Not Running"
@pytest.mark.asyncio
async def test_scaffold_invalid_conversation(client, auth_headers, enable_ai):
"""POST /ai/scaffold returns 404 for nonexistent conversation."""
response = await client.post(
"/api/v1/ai/scaffold",
json={"conversation_id": "00000000-0000-0000-0000-000000000000"},
headers=auth_headers,
)
assert response.status_code == 404
# ── Branch detail endpoint ──
@pytest.mark.asyncio
async def test_branch_detail_success(client, auth_headers, enable_ai):
"""POST /ai/branch-detail returns AI-generated branch nodes."""
# Create and scaffold first
start_resp = await client.post(
"/api/v1/ai/start",
json={
"flow_type": "troubleshooting",
"name": "Service Issues",
"description": "Service troubleshooting",
},
headers=auth_headers,
)
conversation_id = start_resp.json()["conversation_id"]
scaffold_mock = _mock_anthropic_response(SCAFFOLD_RESPONSE_JSON)
with patch("app.core.ai_tree_generator_service._get_client") as mock_client:
mock_client.return_value.messages.create = AsyncMock(return_value=scaffold_mock)
await client.post(
"/api/v1/ai/scaffold",
json={"conversation_id": conversation_id},
headers=auth_headers,
)
# Now generate branch detail
detail_mock = _mock_anthropic_response(BRANCH_DETAIL_JSON)
with patch("app.core.ai_tree_generator_service._get_client") as mock_client:
mock_client.return_value.messages.create = AsyncMock(return_value=detail_mock)
response = await client.post(
"/api/v1/ai/branch-detail",
json={
"conversation_id": conversation_id,
"branch_name": "Service Not Running",
},
headers=auth_headers,
)
assert response.status_code == 200
data = response.json()
assert data["branch_name"] == "Service Not Running"
assert data["steps"]["id"] == "svc-root"
assert data["steps"]["type"] == "decision"
# ── Assemble endpoint ──
@pytest.mark.asyncio
async def test_assemble_success(client, auth_headers, enable_ai):
"""POST /ai/assemble returns assembled tree from branches with detail."""
# Create conversation
start_resp = await client.post(
"/api/v1/ai/start",
json={
"flow_type": "troubleshooting",
"name": "Service Issues",
"description": "Service troubleshooting",
},
headers=auth_headers,
)
conversation_id = start_resp.json()["conversation_id"]
# Scaffold
scaffold_mock = _mock_anthropic_response(SCAFFOLD_RESPONSE_JSON)
with patch("app.core.ai_tree_generator_service._get_client") as mock_client:
mock_client.return_value.messages.create = AsyncMock(return_value=scaffold_mock)
await client.post(
"/api/v1/ai/scaffold",
json={"conversation_id": conversation_id},
headers=auth_headers,
)
# Assemble with branch detail included
branch_tree = json.loads(BRANCH_DETAIL_JSON)
response = await client.post(
"/api/v1/ai/assemble",
json={
"conversation_id": conversation_id,
"selected_branches": [
{
"name": "Service Not Running",
"description": "The target service is stopped.",
"steps": branch_tree,
},
{
"name": "Authentication Failures",
"description": "Users cannot authenticate.",
"steps": branch_tree,
},
],
},
headers=auth_headers,
)
assert response.status_code == 200
data = response.json()
assert data["status"] == "completed"
assert data["suggested_name"] == "Service Issues"
assert "tree_structure" in data
assert data["tree_structure"]["type"] == "decision"
assert data["summary"]["node_count"] > 0
assert data["summary"]["solution_count"] >= 2
@pytest.mark.asyncio
async def test_assemble_requires_min_2_branches(client, auth_headers, enable_ai):
"""POST /ai/assemble rejects fewer than 2 branches."""
start_resp = await client.post(
"/api/v1/ai/start",
json={
"flow_type": "troubleshooting",
"name": "Test",
"description": "Test",
},
headers=auth_headers,
)
conversation_id = start_resp.json()["conversation_id"]
response = await client.post(
"/api/v1/ai/assemble",
json={
"conversation_id": conversation_id,
"selected_branches": [
{"name": "Only Branch", "description": "Just one"},
],
},
headers=auth_headers,
)
assert response.status_code == 422