mirror of
https://github.com/home-assistant/core.git
synced 2026-04-17 23:53:49 +01:00
Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> Co-authored-by: Joostlek <joostlek@outlook.com>
886 lines
28 KiB
Python
886 lines
28 KiB
Python
"""Test the Ollama config flow."""
|
|
|
|
import asyncio
|
|
from unittest.mock import ANY, AsyncMock, patch
|
|
|
|
from httpx import ConnectError
|
|
from ollama import ResponseError
|
|
import pytest
|
|
|
|
from homeassistant import config_entries
|
|
from homeassistant.components import ollama
|
|
from homeassistant.components.ollama.const import DOMAIN
|
|
from homeassistant.config_entries import SOURCE_USER
|
|
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(DOMAIN)
|
|
MockConfigEntry(
|
|
domain=DOMAIN,
|
|
state=config_entries.ConfigEntryState.LOADED,
|
|
).add_to_hass(hass)
|
|
|
|
result = await hass.config_entries.flow.async_init(
|
|
DOMAIN, context={"source": 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"}
|
|
|
|
# No subentries created by default
|
|
assert len(result2.get("subentries", [])) == 0
|
|
assert len(mock_setup_entry.mock_calls) == 1
|
|
assert CONF_API_KEY not in result2["data"]
|
|
|
|
|
|
async def test_duplicate_entry(hass: HomeAssistant) -> None:
|
|
"""Test we abort on duplicate config entry."""
|
|
MockConfigEntry(
|
|
domain=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(
|
|
DOMAIN, context={"source": 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": 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": 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": 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": 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(
|
|
("init_data", "input_data", "expected_data"),
|
|
[
|
|
(
|
|
{
|
|
CONF_URL: "http://localhost:11434",
|
|
CONF_API_KEY: "old-api-key",
|
|
},
|
|
{
|
|
CONF_API_KEY: "new-api-key",
|
|
},
|
|
{
|
|
CONF_URL: "http://localhost:11434",
|
|
CONF_API_KEY: "new-api-key",
|
|
},
|
|
),
|
|
(
|
|
{
|
|
CONF_URL: "http://localhost:11434",
|
|
CONF_API_KEY: "old-api-key",
|
|
},
|
|
{
|
|
# Reconfigure without api_key to test that it gets removed from data
|
|
},
|
|
{
|
|
CONF_URL: "http://localhost:11434",
|
|
},
|
|
),
|
|
],
|
|
)
|
|
async def test_reauth_flow_success(
|
|
hass: HomeAssistant, init_data, input_data, expected_data
|
|
) -> None:
|
|
"""Test successful reauthentication flow."""
|
|
entry = MockConfigEntry(
|
|
domain=DOMAIN,
|
|
data=init_data,
|
|
options={CONF_API_KEY: "stale-options-api-key"},
|
|
version=3,
|
|
minor_version=3,
|
|
)
|
|
entry.add_to_hass(hass)
|
|
|
|
result = await entry.start_reauth_flow(hass)
|
|
assert result["type"] is FlowResultType.FORM
|
|
assert result["step_id"] == "reauth_confirm"
|
|
|
|
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"],
|
|
input_data,
|
|
)
|
|
await hass.async_block_till_done()
|
|
|
|
assert result["type"] is FlowResultType.ABORT
|
|
assert result["reason"] == "reauth_successful"
|
|
|
|
assert entry.data == expected_data
|
|
assert entry.options == {}
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("side_effect", "error"),
|
|
[
|
|
(ResponseError(error="Unauthorized", status_code=401), "invalid_auth"),
|
|
(ConnectError(message="Connection failed"), "cannot_connect"),
|
|
],
|
|
)
|
|
async def test_reauth_flow_errors(hass: HomeAssistant, side_effect, error) -> None:
|
|
"""Test reauthentication flow when authentication fails."""
|
|
entry = MockConfigEntry(
|
|
domain=DOMAIN,
|
|
data={
|
|
CONF_URL: "http://localhost:11434",
|
|
CONF_API_KEY: "old-api-key",
|
|
},
|
|
version=3,
|
|
minor_version=3,
|
|
)
|
|
entry.add_to_hass(hass)
|
|
|
|
result = await entry.start_reauth_flow(hass)
|
|
assert result["type"] is FlowResultType.FORM
|
|
assert result["step_id"] == "reauth_confirm"
|
|
|
|
with patch(
|
|
"homeassistant.components.ollama.config_flow.ollama.AsyncClient.list",
|
|
side_effect=side_effect,
|
|
):
|
|
result = await hass.config_entries.flow.async_configure(
|
|
result["flow_id"],
|
|
{
|
|
CONF_API_KEY: "other-api-key",
|
|
},
|
|
)
|
|
|
|
assert result["type"] is FlowResultType.FORM
|
|
assert result["step_id"] == "reauth_confirm"
|
|
assert result["errors"] == {"base": error}
|
|
|
|
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"],
|
|
{
|
|
CONF_API_KEY: "new-api-key",
|
|
},
|
|
)
|
|
await hass.async_block_till_done()
|
|
|
|
assert result["type"] is FlowResultType.ABORT
|
|
assert result["reason"] == "reauth_successful"
|
|
|
|
assert entry.data == {
|
|
CONF_URL: "http://localhost:11434",
|
|
CONF_API_KEY: "new-api-key",
|
|
}
|
|
|
|
|
|
@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(
|
|
DOMAIN, context={"source": 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}
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("side_effect", "error"),
|
|
[
|
|
(ConnectError(message=""), "cannot_connect"),
|
|
(RuntimeError(), "unknown"),
|
|
],
|
|
)
|
|
async def test_form_errors_recovery(hass: HomeAssistant, side_effect, error) -> None:
|
|
"""Test that the user flow recovers after an error and completes successfully."""
|
|
result = await hass.config_entries.flow.async_init(
|
|
DOMAIN, context={"source": SOURCE_USER}
|
|
)
|
|
|
|
# First attempt fails
|
|
with patch(
|
|
"homeassistant.components.ollama.config_flow.ollama.AsyncClient.list",
|
|
side_effect=side_effect,
|
|
):
|
|
result = await hass.config_entries.flow.async_configure(
|
|
result["flow_id"], {ollama.CONF_URL: "http://localhost:11434"}
|
|
)
|
|
|
|
assert result["type"] is FlowResultType.FORM
|
|
assert result["errors"] == {"base": error}
|
|
|
|
# Second attempt succeeds
|
|
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"},
|
|
)
|
|
await hass.async_block_till_done()
|
|
|
|
assert result["type"] is FlowResultType.CREATE_ENTRY
|
|
assert result["data"] == {ollama.CONF_URL: "http://localhost:11434"}
|
|
|
|
|
|
async def test_form_invalid_url(hass: HomeAssistant) -> None:
|
|
"""Test we handle invalid URL."""
|
|
result = await hass.config_entries.flow.async_init(
|
|
DOMAIN, context={"source": 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": 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": 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": 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": SOURCE_USER},
|
|
)
|
|
|
|
assert result.get("type") is FlowResultType.ABORT
|
|
assert result.get("reason") == "entry_not_loaded"
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("user_input", "expected_headers", "expected_data"),
|
|
[
|
|
(
|
|
{CONF_URL: "http://localhost:11434", CONF_API_KEY: "my-secret-token"},
|
|
{"Authorization": "Bearer my-secret-token"},
|
|
{CONF_URL: "http://localhost:11434", CONF_API_KEY: "my-secret-token"},
|
|
),
|
|
(
|
|
{CONF_URL: "http://localhost:11434", CONF_API_KEY: ""},
|
|
None,
|
|
{CONF_URL: "http://localhost:11434"},
|
|
),
|
|
(
|
|
{CONF_URL: "http://localhost:11434", CONF_API_KEY: " "},
|
|
None,
|
|
{CONF_URL: "http://localhost:11434"},
|
|
),
|
|
(
|
|
{CONF_URL: "http://localhost:11434"},
|
|
None,
|
|
{CONF_URL: "http://localhost:11434"},
|
|
),
|
|
],
|
|
)
|
|
async def test_user_step_async_client_headers(
|
|
hass: HomeAssistant,
|
|
user_input: dict[str, str],
|
|
expected_headers: dict[str, str] | None,
|
|
expected_data: dict[str, str],
|
|
) -> 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": 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
|
|
assert result["data"] == expected_data
|
|
mock_async_client.assert_called_with(
|
|
host="http://localhost:11434",
|
|
headers=expected_headers,
|
|
verify=ANY,
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("status_code", "error", "error_message", "user_input"),
|
|
[
|
|
(
|
|
400,
|
|
"unknown",
|
|
"Bad Request",
|
|
{
|
|
CONF_URL: "http://localhost:11434",
|
|
CONF_API_KEY: "my-secret-token",
|
|
},
|
|
),
|
|
(
|
|
401,
|
|
"invalid_auth",
|
|
"Unauthorized",
|
|
{
|
|
CONF_URL: "http://localhost:11434",
|
|
CONF_API_KEY: "my-secret-token",
|
|
},
|
|
),
|
|
(
|
|
403,
|
|
"invalid_auth",
|
|
"Unauthorized",
|
|
{
|
|
CONF_URL: "http://localhost:11434",
|
|
CONF_API_KEY: "my-secret-token",
|
|
},
|
|
),
|
|
(
|
|
403,
|
|
"invalid_auth",
|
|
"Forbidden",
|
|
{
|
|
CONF_URL: "http://localhost:11434",
|
|
},
|
|
),
|
|
],
|
|
)
|
|
async def test_user_step_errors(
|
|
hass: HomeAssistant,
|
|
status_code: int,
|
|
error: str,
|
|
error_message: str,
|
|
user_input: dict[str, str],
|
|
) -> None:
|
|
"""Test error handling when ollama returns HTTP 4xx."""
|
|
with patch(
|
|
"homeassistant.components.ollama.config_flow.ollama.AsyncClient"
|
|
) as mock_async_client:
|
|
mock_client_instance = AsyncMock()
|
|
mock_async_client.return_value = mock_client_instance
|
|
|
|
mock_client_instance.list.side_effect = ResponseError(
|
|
error=error_message, status_code=status_code
|
|
)
|
|
|
|
result = await hass.config_entries.flow.async_init(
|
|
DOMAIN, context={"source": 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.FORM
|
|
assert result.get("errors") == {"base": error}
|
|
|
|
|
|
async def test_user_step_trim_url(hass: HomeAssistant) -> None:
|
|
"""Test URL is trimmed before validation and persistence."""
|
|
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": SOURCE_USER}
|
|
)
|
|
assert result["type"] is FlowResultType.FORM
|
|
|
|
result = await hass.config_entries.flow.async_configure(
|
|
result["flow_id"],
|
|
user_input={
|
|
CONF_URL: " http://localhost:11434 ",
|
|
},
|
|
)
|
|
await hass.async_block_till_done()
|
|
|
|
assert result["type"] is FlowResultType.CREATE_ENTRY
|
|
assert result["data"] == {CONF_URL: "http://localhost:11434"}
|
|
mock_async_client.assert_called_with(
|
|
host="http://localhost:11434",
|
|
headers=None,
|
|
verify=ANY,
|
|
)
|