resolutionflow/backend/tests/test_ai_fix_endpoint.py

"""Integration tests for the POST /ai/fix-tree endpoint."""

import json
from unittest.mock import AsyncMock, MagicMock, patch

import pytest

from app.core.config import settings


# ── Sample tree (has a decision node with only 1 option + 1 child) ──

SAMPLE_TREE = {
    "id": "root",
    "type": "decision",
    "question": "Is the server up?",
    "options": [
        {"id": "opt-yes", "label": "Yes", "next_node_id": "check-logs"},
        {"id": "opt-no", "label": "No", "next_node_id": "restart"},
    ],
    "children": [
        {
            "id": "check-logs",
            "type": "action",
            "title": "Check Logs",
            "description": "Review logs.",
            "next_node_id": "logs-ok",
        },
        {
            "id": "logs-ok",
            "type": "solution",
            "title": "Logs OK",
            "description": "Issue in logs.",
        },
        {
            "id": "restart",
            "type": "decision",
            "question": "Did restart work?",
            "options": [{"id": "opt-r", "label": "Yes", "next_node_id": "done"}],
            "children": [
                {
                    "id": "done",
                    "type": "solution",
                    "title": "Done",
                    "description": "Fixed.",
                }
            ],
        },
    ],
}

# Fixed version of the "restart" node — 2 options, 2 children
FIXED_RESTART_NODE = {
    "id": "restart",
    "type": "decision",
    "question": "Did restart work?",
    "options": [
        {"id": "opt-r-yes", "label": "Yes", "next_node_id": "done"},
        {"id": "opt-r-no", "label": "No", "next_node_id": "escalate"},
    ],
    "children": [
        {
            "id": "done",
            "type": "solution",
            "title": "Done",
            "description": "Fixed.",
        },
        {
            "id": "escalate",
            "type": "solution",
            "title": "Escalate",
            "description": "Escalate to senior engineer.",
        },
    ],
}

FIX_REQUEST_BODY = {
    "tree_structure": SAMPLE_TREE,
    "tree_name": "Server Troubleshooting",
    "tree_type": "troubleshooting",
    "validation_errors": [
        {
            "node_id": "restart",
            "message": "Decision node 'restart' must have at least 2 options",
        }
    ],
}


def _mock_ai_provider(response_text: str, input_tokens: int = 50, output_tokens: int = 100):
    """Create a mock provider whose generate_json returns given text."""
    provider = MagicMock()
    provider.generate_json = AsyncMock(return_value=(response_text, input_tokens, output_tokens))
    return provider


@pytest.fixture
def enable_ai():
    """Temporarily enable AI by setting a fake API key."""
    original = settings.GOOGLE_AI_API_KEY
    settings.GOOGLE_AI_API_KEY = "test-key-fake"
    yield
    settings.GOOGLE_AI_API_KEY = original


@pytest.fixture
def disable_ai():
    """Ensure AI is disabled."""
    orig_google = settings.GOOGLE_AI_API_KEY
    orig_anthropic = settings.ANTHROPIC_API_KEY
    settings.GOOGLE_AI_API_KEY = None
    settings.ANTHROPIC_API_KEY = None
    yield
    settings.GOOGLE_AI_API_KEY = orig_google
    settings.ANTHROPIC_API_KEY = orig_anthropic


# ── Tests ──


@pytest.mark.asyncio
async def test_returns_401_without_auth(client):
    """POST /ai/fix-tree without auth token returns 401."""
    response = await client.post("/api/v1/ai/fix-tree", json=FIX_REQUEST_BODY)
    assert response.status_code == 401


@pytest.mark.asyncio
async def test_returns_503_when_ai_disabled(client, auth_headers, disable_ai):
    """POST /ai/fix-tree returns 503 when no AI keys are configured."""
    response = await client.post(
        "/api/v1/ai/fix-tree",
        json=FIX_REQUEST_BODY,
        headers=auth_headers,
    )
    assert response.status_code == 503


@pytest.mark.asyncio
async def test_returns_fixes_on_success(client, auth_headers, enable_ai):
    """POST /ai/fix-tree returns fix proposals when AI succeeds."""
    mock_provider = _mock_ai_provider(json.dumps(FIXED_RESTART_NODE))

    with patch(
        "app.core.ai_fix_service.get_ai_provider",
        return_value=mock_provider,
    ):
        response = await client.post(
            "/api/v1/ai/fix-tree",
            json=FIX_REQUEST_BODY,
            headers=auth_headers,
        )

    assert response.status_code == 200
    data = response.json()
    assert "fixes" in data
    assert "tokens_used" in data
    assert len(data["fixes"]) == 1

    fix = data["fixes"][0]
    assert fix["target_node_id"] == "restart"
    assert fix["error_message"] == "Decision node 'restart' must have at least 2 options"
    assert fix["original_node"]["id"] == "restart"
    assert fix["fixed_node"]["id"] == "restart"
    assert len(fix["fixed_node"]["options"]) == 2
    assert len(fix["fixed_node"]["children"]) == 2

    assert data["tokens_used"]["input"] == 50
    assert data["tokens_used"]["output"] == 100