test_config_flow.py

"""Test the Ollama config flow."""

import asyncio
from unittest.mock import ANY, AsyncMock, patch

from httpx import ConnectError
import pytest

from homeassistant import config_entries
from homeassistant.components import ollama
from homeassistant.components.ollama.const import DOMAIN
from homeassistant.const import CONF_API_KEY, CONF_LLM_HASS_API, CONF_NAME, CONF_URL
from homeassistant.core import HomeAssistant
from homeassistant.data_entry_flow import FlowResultType

from tests.common import MockConfigEntry

TEST_MODEL = "test_model:latest"


async def test_form(hass: HomeAssistant) -> None:
    """Test flow when configuring URL only."""
    # Pretend we already set up a config entry.
    hass.config.components.add(ollama.DOMAIN)
    MockConfigEntry(
        domain=ollama.DOMAIN,
        state=config_entries.ConfigEntryState.LOADED,
    ).add_to_hass(hass)

    result = await hass.config_entries.flow.async_init(
        ollama.DOMAIN, context={"source": config_entries.SOURCE_USER}
    )
    assert result["type"] is FlowResultType.FORM
    assert result["errors"] is None

    with (
        patch(
            "homeassistant.components.ollama.config_flow.ollama.AsyncClient.list",
            return_value={"models": [{"model": TEST_MODEL}]},
        ),
        patch(
            "homeassistant.components.ollama.async_setup_entry",
            return_value=True,
        ) as mock_setup_entry,
    ):
        result2 = await hass.config_entries.flow.async_configure(
            result["flow_id"], {ollama.CONF_URL: "http://localhost:11434"}
        )
        await hass.async_block_till_done()

    assert result2["type"] is FlowResultType.CREATE_ENTRY
    assert result2["data"] == {
        ollama.CONF_URL: "http://localhost:11434",
        ollama.CONF_API_KEY: "",  # Default API key should be empty string
    }
    # No subentries created by default
    assert len(result2.get("subentries", [])) == 0
    assert len(mock_setup_entry.mock_calls) == 1


async def test_duplicate_entry(hass: HomeAssistant) -> None:
    """Test we abort on duplicate config entry."""
    MockConfigEntry(
        domain=ollama.DOMAIN,
        data={
            ollama.CONF_URL: "http://localhost:11434",
            ollama.CONF_MODEL: "test_model",
        },
    ).add_to_hass(hass)

    result = await hass.config_entries.flow.async_init(
        ollama.DOMAIN, context={"source": config_entries.SOURCE_USER}
    )
    assert result["type"] is FlowResultType.FORM
    assert not result["errors"]

    with patch(
        "homeassistant.components.ollama.config_flow.ollama.AsyncClient.list",
        return_value={"models": [{"model": "test_model"}]},
    ):
        result = await hass.config_entries.flow.async_configure(
            result["flow_id"],
            {
                ollama.CONF_URL: "http://localhost:11434",
            },
        )

    assert result["type"] is FlowResultType.ABORT
    assert result["reason"] == "already_configured"


async def test_subentry_options(
    hass: HomeAssistant, mock_config_entry, mock_init_component
) -> None:
    """Test the subentry options form."""
    subentry = next(iter(mock_config_entry.subentries.values()))

    # Test reconfiguration
    with patch(
        "ollama.AsyncClient.list",
        return_value={"models": [{"model": TEST_MODEL}]},
    ):
        options_flow = await mock_config_entry.start_subentry_reconfigure_flow(
            hass, subentry.subentry_id
        )

        assert options_flow["type"] is FlowResultType.FORM
        assert options_flow["step_id"] == "set_options"

        options = await hass.config_entries.subentries.async_configure(
            options_flow["flow_id"],
            {
                ollama.CONF_MODEL: TEST_MODEL,
                ollama.CONF_PROMPT: "test prompt",
                ollama.CONF_MAX_HISTORY: 100,
                ollama.CONF_NUM_CTX: 32768,
                ollama.CONF_THINK: True,
            },
        )
    await hass.async_block_till_done()

    assert options["type"] is FlowResultType.ABORT
    assert options["reason"] == "reconfigure_successful"
    assert subentry.data == {
        ollama.CONF_MODEL: TEST_MODEL,
        ollama.CONF_PROMPT: "test prompt",
        ollama.CONF_MAX_HISTORY: 100.0,
        ollama.CONF_NUM_CTX: 32768.0,
        ollama.CONF_THINK: True,
    }


async def test_creating_new_conversation_subentry(
    hass: HomeAssistant,
    mock_config_entry: MockConfigEntry,
    mock_init_component,
) -> None:
    """Test creating a new conversation subentry includes name field."""
    # Start a new subentry flow
    with patch(
        "ollama.AsyncClient.list",
        return_value={"models": [{"model": TEST_MODEL}]},
    ):
        new_flow = await hass.config_entries.subentries.async_init(
            (mock_config_entry.entry_id, "conversation"),
            context={"source": config_entries.SOURCE_USER},
        )

        assert new_flow["type"] is FlowResultType.FORM
        assert new_flow["step_id"] == "set_options"

        # Configure the new subentry with name field
        result = await hass.config_entries.subentries.async_configure(
            new_flow["flow_id"],
            {
                ollama.CONF_MODEL: TEST_MODEL,
                CONF_NAME: "New Test Conversation",
                ollama.CONF_PROMPT: "new test prompt",
                ollama.CONF_MAX_HISTORY: 50,
                ollama.CONF_NUM_CTX: 16384,
                ollama.CONF_THINK: False,
            },
        )
    await hass.async_block_till_done()

    assert result["type"] is FlowResultType.CREATE_ENTRY
    assert result["title"] == "New Test Conversation"
    assert result["data"] == {
        ollama.CONF_MODEL: TEST_MODEL,
        ollama.CONF_PROMPT: "new test prompt",
        ollama.CONF_MAX_HISTORY: 50.0,
        ollama.CONF_NUM_CTX: 16384.0,
        ollama.CONF_THINK: False,
    }


async def test_creating_conversation_subentry_not_loaded(
    hass: HomeAssistant,
    mock_init_component,
    mock_config_entry: MockConfigEntry,
) -> None:
    """Test creating a conversation subentry when entry is not loaded."""
    await hass.config_entries.async_unload(mock_config_entry.entry_id)
    result = await hass.config_entries.subentries.async_init(
        (mock_config_entry.entry_id, "conversation"),
        context={"source": config_entries.SOURCE_USER},
    )

    assert result["type"] is FlowResultType.ABORT
    assert result["reason"] == "entry_not_loaded"


async def test_subentry_need_download(
    hass: HomeAssistant,
    mock_init_component,
    mock_config_entry: MockConfigEntry,
) -> None:
    """Test subentry creation when model needs to be downloaded."""

    async def delayed_pull(self, model: str) -> None:
        """Simulate a delayed model download."""
        assert model == "llama3.2:latest"
        await asyncio.sleep(0)  # yield the event loop 1 iteration

    with (
        patch(
            "ollama.AsyncClient.list",
            return_value={"models": [{"model": TEST_MODEL}]},
        ),
        patch("ollama.AsyncClient.pull", delayed_pull),
    ):
        new_flow = await hass.config_entries.subentries.async_init(
            (mock_config_entry.entry_id, "conversation"),
            context={"source": config_entries.SOURCE_USER},
        )

        assert new_flow["type"] is FlowResultType.FORM, new_flow
        assert new_flow["step_id"] == "set_options"

        # Configure the new subentry with a model that needs downloading
        result = await hass.config_entries.subentries.async_configure(
            new_flow["flow_id"],
            {
                ollama.CONF_MODEL: "llama3.2:latest",  # not cached
                CONF_NAME: "New Test Conversation",
                ollama.CONF_PROMPT: "new test prompt",
                ollama.CONF_MAX_HISTORY: 50,
                ollama.CONF_NUM_CTX: 16384,
                ollama.CONF_THINK: False,
            },
        )

        assert result["type"] is FlowResultType.SHOW_PROGRESS
        assert result["step_id"] == "download"
        assert result["progress_action"] == "download"

        await hass.async_block_till_done()

        result = await hass.config_entries.subentries.async_configure(
            new_flow["flow_id"], {}
        )

    assert result["type"] is FlowResultType.CREATE_ENTRY
    assert result["title"] == "New Test Conversation"
    assert result["data"] == {
        ollama.CONF_MODEL: "llama3.2:latest",
        ollama.CONF_PROMPT: "new test prompt",
        ollama.CONF_MAX_HISTORY: 50.0,
        ollama.CONF_NUM_CTX: 16384.0,
        ollama.CONF_THINK: False,
    }


async def test_subentry_download_error(
    hass: HomeAssistant,
    mock_init_component,
    mock_config_entry: MockConfigEntry,
) -> None:
    """Test subentry creation when model download fails."""

    async def delayed_pull(self, model: str) -> None:
        """Simulate a delayed model download."""
        await asyncio.sleep(0)  # yield

        raise RuntimeError("Download failed")

    with (
        patch(
            "ollama.AsyncClient.list",
            return_value={"models": [{"model": TEST_MODEL}]},
        ),
        patch("ollama.AsyncClient.pull", delayed_pull),
    ):
        new_flow = await hass.config_entries.subentries.async_init(
            (mock_config_entry.entry_id, "conversation"),
            context={"source": config_entries.SOURCE_USER},
        )

        assert new_flow["type"] is FlowResultType.FORM
        assert new_flow["step_id"] == "set_options"

        # Configure with a model that needs downloading but will fail
        result = await hass.config_entries.subentries.async_configure(
            new_flow["flow_id"],
            {
                ollama.CONF_MODEL: "llama3.2:latest",
                CONF_NAME: "New Test Conversation",
                ollama.CONF_PROMPT: "new test prompt",
                ollama.CONF_MAX_HISTORY: 50,
                ollama.CONF_NUM_CTX: 16384,
                ollama.CONF_THINK: False,
            },
        )

        # Should show progress flow result for download
        assert result["type"] is FlowResultType.SHOW_PROGRESS
        assert result["step_id"] == "download"
        assert result["progress_action"] == "download"

        # Wait for download task to complete (with error)
        await hass.async_block_till_done()

        # Submit the progress flow - should get failure
        result = await hass.config_entries.subentries.async_configure(
            new_flow["flow_id"], {}
        )

    assert result["type"] is FlowResultType.ABORT
    assert result["reason"] == "download_failed"


@pytest.mark.parametrize(
    ("side_effect", "error"),
    [
        (ConnectError(message=""), "cannot_connect"),
        (RuntimeError(), "unknown"),
    ],
)
async def test_form_errors(hass: HomeAssistant, side_effect, error) -> None:
    """Test we handle errors."""
    result = await hass.config_entries.flow.async_init(
        ollama.DOMAIN, context={"source": config_entries.SOURCE_USER}
    )

    with patch(
        "homeassistant.components.ollama.config_flow.ollama.AsyncClient.list",
        side_effect=side_effect,
    ):
        result2 = await hass.config_entries.flow.async_configure(
            result["flow_id"], {ollama.CONF_URL: "http://localhost:11434"}
        )

    assert result2["type"] is FlowResultType.FORM
    assert result2["errors"] == {"base": error}


async def test_form_invalid_url(hass: HomeAssistant) -> None:
    """Test we handle invalid URL."""
    result = await hass.config_entries.flow.async_init(
        ollama.DOMAIN, context={"source": config_entries.SOURCE_USER}
    )

    result2 = await hass.config_entries.flow.async_configure(
        result["flow_id"], {ollama.CONF_URL: "not-a-valid-url"}
    )

    assert result2["type"] is FlowResultType.FORM
    assert result2["errors"] == {"base": "invalid_url"}


async def test_subentry_connection_error(
    hass: HomeAssistant,
    mock_init_component,
    mock_config_entry: MockConfigEntry,
) -> None:
    """Test subentry creation when connection to Ollama server fails."""
    with patch(
        "ollama.AsyncClient.list",
        side_effect=ConnectError("Connection failed"),
    ):
        new_flow = await hass.config_entries.subentries.async_init(
            (mock_config_entry.entry_id, "conversation"),
            context={"source": config_entries.SOURCE_USER},
        )

    assert new_flow["type"] is FlowResultType.ABORT
    assert new_flow["reason"] == "cannot_connect"


async def test_subentry_model_check_exception(
    hass: HomeAssistant,
    mock_init_component,
    mock_config_entry: MockConfigEntry,
) -> None:
    """Test subentry creation when checking model availability throws exception."""
    with patch(
        "ollama.AsyncClient.list",
        side_effect=[
            {"models": [{"model": TEST_MODEL}]},  # First call succeeds
            RuntimeError("Failed to check models"),  # Second call fails
        ],
    ):
        new_flow = await hass.config_entries.subentries.async_init(
            (mock_config_entry.entry_id, "conversation"),
            context={"source": config_entries.SOURCE_USER},
        )

        assert new_flow["type"] is FlowResultType.FORM
        assert new_flow["step_id"] == "set_options"

        # Configure with a model, should fail when checking availability
        result = await hass.config_entries.subentries.async_configure(
            new_flow["flow_id"],
            {
                ollama.CONF_MODEL: "new_model:latest",
                CONF_NAME: "Test Conversation",
                ollama.CONF_PROMPT: "test prompt",
                ollama.CONF_MAX_HISTORY: 50,
                ollama.CONF_NUM_CTX: 16384,
                ollama.CONF_THINK: False,
            },
        )

    assert result["type"] is FlowResultType.ABORT
    assert result["reason"] == "cannot_connect"


async def test_subentry_reconfigure_with_download(
    hass: HomeAssistant,
    mock_init_component,
    mock_config_entry: MockConfigEntry,
) -> None:
    """Test reconfiguring subentry when model needs to be downloaded."""
    subentry = next(iter(mock_config_entry.subentries.values()))

    async def delayed_pull(self, model: str) -> None:
        """Simulate a delayed model download."""
        assert model == "llama3.2:latest"
        await asyncio.sleep(0)  # yield the event loop

    with (
        patch(
            "ollama.AsyncClient.list",
            return_value={"models": [{"model": TEST_MODEL}]},
        ),
        patch("ollama.AsyncClient.pull", delayed_pull),
    ):
        reconfigure_flow = await mock_config_entry.start_subentry_reconfigure_flow(
            hass, subentry.subentry_id
        )

        assert reconfigure_flow["type"] is FlowResultType.FORM
        assert reconfigure_flow["step_id"] == "set_options"

        # Reconfigure with a model that needs downloading
        result = await hass.config_entries.subentries.async_configure(
            reconfigure_flow["flow_id"],
            {
                ollama.CONF_MODEL: "llama3.2:latest",
                ollama.CONF_PROMPT: "updated prompt",
                ollama.CONF_MAX_HISTORY: 75,
                ollama.CONF_NUM_CTX: 8192,
                ollama.CONF_THINK: True,
            },
        )

        assert result["type"] is FlowResultType.SHOW_PROGRESS
        assert result["step_id"] == "download"

        await hass.async_block_till_done()

        # Finish download
        result = await hass.config_entries.subentries.async_configure(
            reconfigure_flow["flow_id"], {}
        )

    assert result["type"] is FlowResultType.ABORT
    assert result["reason"] == "reconfigure_successful"
    assert subentry.data == {
        ollama.CONF_MODEL: "llama3.2:latest",
        ollama.CONF_PROMPT: "updated prompt",
        ollama.CONF_MAX_HISTORY: 75.0,
        ollama.CONF_NUM_CTX: 8192.0,
        ollama.CONF_THINK: True,
    }


async def test_filter_invalid_llms(
    hass: HomeAssistant,
    mock_init_component,
    mock_config_entry_with_assist_invalid_api: MockConfigEntry,
) -> None:
    """Test reconfiguring subentry when one of the configured LLM APIs has been removed."""
    subentry = next(iter(mock_config_entry_with_assist_invalid_api.subentries.values()))

    assert len(subentry.data.get(CONF_LLM_HASS_API)) == 2
    assert "invalid_api" in subentry.data.get(CONF_LLM_HASS_API)
    assert "assist" in subentry.data.get(CONF_LLM_HASS_API)

    valid_apis = ollama.config_flow.filter_invalid_llm_apis(
        hass, subentry.data[CONF_LLM_HASS_API]
    )

    assert len(valid_apis) == 1
    assert "invalid_api" not in valid_apis
    assert "assist" in valid_apis


async def test_creating_ai_task_subentry(
    hass: HomeAssistant,
    mock_config_entry: MockConfigEntry,
    mock_init_component,
) -> None:
    """Test creating an AI task subentry."""
    old_subentries = set(mock_config_entry.subentries)
    # Original conversation + original ai_task
    assert len(mock_config_entry.subentries) == 2

    with patch(
        "ollama.AsyncClient.list",
        return_value={"models": [{"model": "test_model:latest"}]},
    ):
        result = await hass.config_entries.subentries.async_init(
            (mock_config_entry.entry_id, "ai_task_data"),
            context={"source": config_entries.SOURCE_USER},
        )

    assert result.get("type") is FlowResultType.FORM
    assert result.get("step_id") == "set_options"
    assert not result.get("errors")

    with patch(
        "ollama.AsyncClient.list",
        return_value={"models": [{"model": "test_model:latest"}]},
    ):
        result2 = await hass.config_entries.subentries.async_configure(
            result["flow_id"],
            {
                "name": "Custom AI Task",
                ollama.CONF_MODEL: "test_model:latest",
                ollama.CONF_MAX_HISTORY: 5,
                ollama.CONF_NUM_CTX: 4096,
                ollama.CONF_KEEP_ALIVE: 30,
                ollama.CONF_THINK: False,
            },
        )
        await hass.async_block_till_done()

    assert result2.get("type") is FlowResultType.CREATE_ENTRY
    assert result2.get("title") == "Custom AI Task"
    assert result2.get("data") == {
        ollama.CONF_MODEL: "test_model:latest",
        ollama.CONF_MAX_HISTORY: 5,
        ollama.CONF_NUM_CTX: 4096,
        ollama.CONF_KEEP_ALIVE: 30,
        ollama.CONF_THINK: False,
    }

    assert (
        len(mock_config_entry.subentries) == 3
    )  # Original conversation + original ai_task + new ai_task

    new_subentry_id = list(set(mock_config_entry.subentries) - old_subentries)[0]
    new_subentry = mock_config_entry.subentries[new_subentry_id]
    assert new_subentry.subentry_type == "ai_task_data"
    assert new_subentry.title == "Custom AI Task"


async def test_ai_task_subentry_not_loaded(
    hass: HomeAssistant,
    mock_config_entry: MockConfigEntry,
) -> None:
    """Test creating an AI task subentry when entry is not loaded."""
    # Don't call mock_init_component to simulate not loaded state
    result = await hass.config_entries.subentries.async_init(
        (mock_config_entry.entry_id, "ai_task_data"),
        context={"source": config_entries.SOURCE_USER},
    )

    assert result.get("type") is FlowResultType.ABORT
    assert result.get("reason") == "entry_not_loaded"


@pytest.mark.parametrize(
    ("user_input", "expected_headers"),
    [
        (
            {CONF_URL: "http://localhost:11434", CONF_API_KEY: "my-secret-token"},
            {"Authorization": "Bearer my-secret-token"},
        ),
        (
            {CONF_URL: "http://localhost:11434", CONF_API_KEY: ""},
            None,
        ),
        (
            {CONF_URL: "http://localhost:11434"},
            None,
        ),
    ],
)
async def test_user_step_async_client_headers(
    hass: HomeAssistant,
    user_input: dict[str, str],
    expected_headers: dict[str, str] | None,
) -> None:
    """Test Authorization header passed to AsyncClient with/without api_key."""
    with patch(
        "homeassistant.components.ollama.config_flow.ollama.AsyncClient",
    ) as mock_async_client:
        mock_async_client.return_value.list = AsyncMock(return_value={"models": []})

        result = await hass.config_entries.flow.async_init(
            DOMAIN, context={"source": "user"}
        )
        assert result["type"] is FlowResultType.FORM

        result = await hass.config_entries.flow.async_configure(
            result["flow_id"],
            user_input=user_input,
        )
        await hass.async_block_till_done()

    assert result["type"] is FlowResultType.CREATE_ENTRY
    mock_async_client.assert_called_with(
        host="http://localhost:11434",
        headers=expected_headers,
        verify=ANY,
    )