Install and start Mosquitto MQTT broker add on from MQTT config flow (#124106)

* Opt in to install Mosquitto broker add-on in MQTT config flow

* rephrase

* Tests with supervisor and running add-on

* Complete tests for success flows

* Also set up entry in success flow

* Use realistic names for addon and broker

* Finetuning and fail test cases

* Spelling

* Improve translation strings

* Update addon docstr

Co-authored-by: Erik Montnemery <erik@montnemery.com>

* Raise AddonError if add-on does not start

* Only show the option to use the add-on

* Simplify flow, rework and cleanup

* Revert unrelated cleanup, process suggestion

* Move ADDON_SLUG const to addon module

* Move fixture to component level

* Move back supervisor fixture

* Move addon_setup_time_fixture and superfixe to config flow model tests

* Refactor hassio fixture

* Rename helpers as they are no fixtures, remove fixture from their names

---------

Co-authored-by: Erik Montnemery <erik@montnemery.com>
This commit is contained in:
Jan Bouwhuis 2024-08-22 09:07:45 +02:00 committed by GitHub
parent 9b9e98a26e
commit 61114d8328
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
7 changed files with 876 additions and 18 deletions

View file

@ -2,8 +2,10 @@
from __future__ import annotations
import asyncio
from collections import OrderedDict
from collections.abc import Callable, Mapping
import logging
import queue
from ssl import PROTOCOL_TLS_CLIENT, SSLContext, SSLError
from types import MappingProxyType
@ -14,7 +16,12 @@ from cryptography.x509 import load_pem_x509_certificate
import voluptuous as vol
from homeassistant.components.file_upload import process_uploaded_file
from homeassistant.components.hassio import HassioServiceInfo
from homeassistant.components.hassio import HassioServiceInfo, is_hassio
from homeassistant.components.hassio.addon_manager import (
AddonError,
AddonManager,
AddonState,
)
from homeassistant.config_entries import (
ConfigEntry,
ConfigFlow,
@ -32,6 +39,7 @@ from homeassistant.const import (
CONF_USERNAME,
)
from homeassistant.core import callback
from homeassistant.data_entry_flow import AbortFlow
from homeassistant.helpers import config_validation as cv
from homeassistant.helpers.json import json_dumps
from homeassistant.helpers.selector import (
@ -51,6 +59,7 @@ from homeassistant.helpers.selector import (
)
from homeassistant.util.json import JSON_DECODE_EXCEPTIONS, json_loads
from .addon import get_addon_manager
from .client import MqttClientSetup
from .const import (
ATTR_PAYLOAD,
@ -91,6 +100,11 @@ from .util import (
valid_publish_topic,
)
_LOGGER = logging.getLogger(__name__)
ADDON_SETUP_TIMEOUT = 5
ADDON_SETUP_TIMEOUT_ROUNDS = 5
MQTT_TIMEOUT = 5
ADVANCED_OPTIONS = "advanced_options"
@ -197,6 +211,12 @@ class FlowHandler(ConfigFlow, domain=DOMAIN):
entry: ConfigEntry | None
_hassio_discovery: dict[str, Any] | None = None
_addon_manager: AddonManager
def __init__(self) -> None:
"""Set up flow instance."""
self.install_task: asyncio.Task | None = None
self.start_task: asyncio.Task | None = None
@staticmethod
@callback
@ -206,6 +226,118 @@ class FlowHandler(ConfigFlow, domain=DOMAIN):
"""Get the options flow for this handler."""
return MQTTOptionsFlowHandler(config_entry)
async def _async_install_addon(self) -> None:
"""Install the Mosquitto Mqtt broker add-on."""
addon_manager: AddonManager = get_addon_manager(self.hass)
await addon_manager.async_schedule_install_addon()
async def async_step_install_failed(
self, user_input: dict[str, Any] | None = None
) -> ConfigFlowResult:
"""Add-on installation failed."""
return self.async_abort(
reason="addon_install_failed",
description_placeholders={"addon": self._addon_manager.addon_name},
)
async def async_step_install_addon(
self, user_input: dict[str, Any] | None = None
) -> ConfigFlowResult:
"""Install Mosquitto Broker add-on."""
if self.install_task is None:
self.install_task = self.hass.async_create_task(self._async_install_addon())
if not self.install_task.done():
return self.async_show_progress(
step_id="install_addon",
progress_action="install_addon",
progress_task=self.install_task,
)
try:
await self.install_task
except AddonError as err:
_LOGGER.error(err)
return self.async_show_progress_done(next_step_id="install_failed")
finally:
self.install_task = None
return self.async_show_progress_done(next_step_id="start_addon")
async def async_step_start_failed(
self, user_input: dict[str, Any] | None = None
) -> ConfigFlowResult:
"""Add-on start failed."""
return self.async_abort(
reason="addon_start_failed",
description_placeholders={"addon": self._addon_manager.addon_name},
)
async def async_step_start_addon(
self, user_input: dict[str, Any] | None = None
) -> ConfigFlowResult:
"""Start Mosquitto Broker add-on."""
if not self.start_task:
self.start_task = self.hass.async_create_task(self._async_start_addon())
if not self.start_task.done():
return self.async_show_progress(
step_id="start_addon",
progress_action="start_addon",
progress_task=self.start_task,
)
try:
await self.start_task
except AddonError as err:
_LOGGER.error(err)
return self.async_show_progress_done(next_step_id="start_failed")
finally:
self.start_task = None
return self.async_show_progress_done(next_step_id="setup_entry_from_discovery")
async def _async_get_config_and_try(self) -> dict[str, Any] | None:
"""Get the MQTT add-on discovery info and try the connection."""
if self._hassio_discovery is not None:
return self._hassio_discovery
addon_manager: AddonManager = get_addon_manager(self.hass)
try:
addon_discovery_config = (
await addon_manager.async_get_addon_discovery_info()
)
config: dict[str, Any] = {
CONF_BROKER: addon_discovery_config[CONF_HOST],
CONF_PORT: addon_discovery_config[CONF_PORT],
CONF_USERNAME: addon_discovery_config.get(CONF_USERNAME),
CONF_PASSWORD: addon_discovery_config.get(CONF_PASSWORD),
CONF_DISCOVERY: DEFAULT_DISCOVERY,
}
except AddonError:
# We do not have discovery information yet
return None
if await self.hass.async_add_executor_job(
try_connection,
config,
):
self._hassio_discovery = config
return config
return None
async def _async_start_addon(self) -> None:
"""Start the Mosquitto Broker add-on."""
addon_manager: AddonManager = get_addon_manager(self.hass)
await addon_manager.async_schedule_start_addon()
# Sleep some seconds to let the add-on start properly before connecting.
for _ in range(ADDON_SETUP_TIMEOUT_ROUNDS):
await asyncio.sleep(ADDON_SETUP_TIMEOUT)
# Finish setup using discovery info to test the connection
if await self._async_get_config_and_try():
break
else:
raise AddonError(
f"Failed to correctly start {addon_manager.addon_name} add-on"
)
async def async_step_user(
self, user_input: dict[str, Any] | None = None
) -> ConfigFlowResult:
@ -213,8 +345,57 @@ class FlowHandler(ConfigFlow, domain=DOMAIN):
if self._async_current_entries():
return self.async_abort(reason="single_instance_allowed")
if is_hassio(self.hass):
# Offer to set up broker add-on if supervisor is available
self._addon_manager = get_addon_manager(self.hass)
return self.async_show_menu(
step_id="user",
menu_options=["addon", "broker"],
description_placeholders={"addon": self._addon_manager.addon_name},
)
# Start up a flow for manual setup
return await self.async_step_broker()
async def async_step_setup_entry_from_discovery(
self, user_input: dict[str, Any] | None = None
) -> ConfigFlowResult:
"""Set up mqtt entry from discovery info."""
if (config := await self._async_get_config_and_try()) is not None:
return self.async_create_entry(
title=self._addon_manager.addon_name,
data=config,
)
raise AbortFlow(
"addon_connection_failed",
description_placeholders={"addon": self._addon_manager.addon_name},
)
async def async_step_addon(
self, user_input: dict[str, Any] | None = None
) -> ConfigFlowResult:
"""Install and start MQTT broker add-on."""
addon_manager = self._addon_manager
try:
addon_info = await addon_manager.async_get_addon_info()
except AddonError as err:
raise AbortFlow(
"addon_info_failed",
description_placeholders={"addon": self._addon_manager.addon_name},
) from err
if addon_info.state == AddonState.RUNNING:
# Finish setup using discovery info
return await self.async_step_setup_entry_from_discovery()
if addon_info.state == AddonState.NOT_RUNNING:
return await self.async_step_start_addon()
# Install the add-on and start it
return await self.async_step_install_addon()
async def async_step_reauth(
self, entry_data: Mapping[str, Any]
) -> ConfigFlowResult:
@ -293,7 +474,7 @@ class FlowHandler(ConfigFlow, domain=DOMAIN):
async def async_step_hassio(
self, discovery_info: HassioServiceInfo
) -> ConfigFlowResult:
"""Receive a Hass.io discovery."""
"""Receive a Hass.io discovery or process setup after addon install."""
await self._async_handle_discovery_without_unique_id()
self._hassio_discovery = discovery_info.config