From ce692afead189ac63a2aa1f5b286caae8581e8e3 Mon Sep 17 00:00:00 2001 From: "J. Nick Koston" Date: Thu, 6 May 2021 09:50:28 -0500 Subject: [PATCH] Add rainmachine discovery (#49970) Co-authored-by: Paulus Schoutsen --- .../components/rainmachine/__init__.py | 56 ++-- .../components/rainmachine/config_flow.py | 143 ++++++++--- homeassistant/components/rainmachine/const.py | 1 - .../components/rainmachine/manifest.json | 11 +- .../components/rainmachine/strings.json | 1 + .../rainmachine/translations/en.json | 1 + homeassistant/generated/zeroconf.py | 6 + .../rainmachine/test_config_flow.py | 239 +++++++++++++++--- 8 files changed, 352 insertions(+), 106 deletions(-) diff --git a/homeassistant/components/rainmachine/__init__.py b/homeassistant/components/rainmachine/__init__.py index 8c72922699e..66fcc8939fb 100644 --- a/homeassistant/components/rainmachine/__init__.py +++ b/homeassistant/components/rainmachine/__init__.py @@ -24,7 +24,9 @@ from homeassistant.helpers.update_coordinator import ( DataUpdateCoordinator, UpdateFailed, ) +from homeassistant.util.network import is_ip_address +from .config_flow import get_client_controller from .const import ( CONF_ZONE_RUN_TIME, DATA_CONTROLLER, @@ -38,8 +40,6 @@ from .const import ( LOGGER, ) -DATA_LISTENER = "listener" - DEFAULT_ATTRIBUTION = "Data provided by Green Electronics LLC" DEFAULT_ICON = "mdi:water" DEFAULT_SSL = True @@ -70,32 +70,10 @@ async def async_update_programs_and_zones( ) -async def async_setup(hass: HomeAssistant, config: dict) -> bool: - """Set up the RainMachine component.""" - hass.data[DOMAIN] = {DATA_CONTROLLER: {}, DATA_COORDINATOR: {}, DATA_LISTENER: {}} - return True - - async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: """Set up RainMachine as config entry.""" + hass.data.setdefault(DOMAIN, {DATA_CONTROLLER: {}, DATA_COORDINATOR: {}}) hass.data[DOMAIN][DATA_COORDINATOR][entry.entry_id] = {} - - entry_updates = {} - if not entry.unique_id: - # If the config entry doesn't already have a unique ID, set one: - entry_updates["unique_id"] = entry.data[CONF_IP_ADDRESS] - if CONF_ZONE_RUN_TIME in entry.data: - # If a zone run time exists in the config entry's data, pop it and move it to - # options: - data = {**entry.data} - entry_updates["data"] = data - entry_updates["options"] = { - **entry.options, - CONF_ZONE_RUN_TIME: data.pop(CONF_ZONE_RUN_TIME), - } - if entry_updates: - hass.config_entries.async_update_entry(entry, **entry_updates) - websession = aiohttp_client.async_get_clientsession(hass) client = Client(session=websession) @@ -107,14 +85,29 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: ssl=entry.data.get(CONF_SSL, DEFAULT_SSL), ) except RainMachineError as err: - LOGGER.error("An error occurred: %s", err) raise ConfigEntryNotReady from err # regenmaschine can load multiple controllers at once, but we only grab the one # we loaded above: - controller = hass.data[DOMAIN][DATA_CONTROLLER][entry.entry_id] = next( - iter(client.controllers.values()) - ) + controller = hass.data[DOMAIN][DATA_CONTROLLER][ + entry.entry_id + ] = get_client_controller(client) + + entry_updates = {} + if not entry.unique_id or is_ip_address(entry.unique_id): + # If the config entry doesn't already have a unique ID, set one: + entry_updates["unique_id"] = controller.mac + if CONF_ZONE_RUN_TIME in entry.data: + # If a zone run time exists in the config entry's data, pop it and move it to + # options: + data = {**entry.data} + entry_updates["data"] = data + entry_updates["options"] = { + **entry.options, + CONF_ZONE_RUN_TIME: data.pop(CONF_ZONE_RUN_TIME), + } + if entry_updates: + hass.config_entries.async_update_entry(entry, **entry_updates) async def async_update(api_category: str) -> dict: """Update the appropriate API data based on a category.""" @@ -158,7 +151,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: hass.config_entries.async_setup_platforms(entry, PLATFORMS) - hass.data[DOMAIN][DATA_LISTENER] = entry.add_update_listener(async_reload_entry) + entry.async_on_unload(entry.add_update_listener(async_reload_entry)) return True @@ -168,9 +161,6 @@ async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: unload_ok = await hass.config_entries.async_unload_platforms(entry, PLATFORMS) if unload_ok: hass.data[DOMAIN][DATA_COORDINATOR].pop(entry.entry_id) - cancel_listener = hass.data[DOMAIN][DATA_LISTENER].pop(entry.entry_id) - cancel_listener() - return unload_ok diff --git a/homeassistant/components/rainmachine/config_flow.py b/homeassistant/components/rainmachine/config_flow.py index 37b7da4b56b..8ef21fb185e 100644 --- a/homeassistant/components/rainmachine/config_flow.py +++ b/homeassistant/components/rainmachine/config_flow.py @@ -6,7 +6,9 @@ import voluptuous as vol from homeassistant import config_entries from homeassistant.const import CONF_IP_ADDRESS, CONF_PASSWORD, CONF_PORT, CONF_SSL from homeassistant.core import callback +from homeassistant.data_entry_flow import AbortFlow from homeassistant.helpers import aiohttp_client, config_validation as cv +from homeassistant.helpers.typing import DiscoveryInfoType from .const import CONF_ZONE_RUN_TIME, DEFAULT_PORT, DEFAULT_ZONE_RUN, DOMAIN @@ -19,58 +21,129 @@ DATA_SCHEMA = vol.Schema( ) +def get_client_controller(client): + """Enumerate controllers to find the first mac.""" + for controller in client.controllers.values(): + return controller + + +async def async_get_controller(hass, ip_address, password, port, ssl): + """Auth and fetch the mac address from the controller.""" + websession = aiohttp_client.async_get_clientsession(hass) + client = Client(session=websession) + try: + await client.load_local(ip_address, password, port=port, ssl=ssl) + except RainMachineError: + return None + else: + return get_client_controller(client) + + class RainMachineFlowHandler(config_entries.ConfigFlow, domain=DOMAIN): """Handle a RainMachine config flow.""" VERSION = 1 + def __init__(self): + """Initialize config flow.""" + self.discovered_ip_address = None + @staticmethod @callback def async_get_options_flow(config_entry): """Define the config flow to handle options.""" return RainMachineOptionsFlowHandler(config_entry) + @callback + def _async_abort_ip_address_configured(self, ip_address): + """Abort if we already have an entry for the ip.""" + # IP already configured + for entry in self._async_current_entries(include_ignore=False): + if ip_address == entry.data[CONF_IP_ADDRESS]: + raise AbortFlow("already_configured") + + async def async_step_homekit(self, discovery_info): + """Handle a flow initialized by homekit discovery.""" + return await self.async_step_zeroconf(discovery_info) + + async def async_step_zeroconf(self, discovery_info: DiscoveryInfoType): + """Handle discovery via zeroconf.""" + ip_address = discovery_info["host"] + + self._async_abort_ip_address_configured(ip_address) + # Handle IP change + for entry in self._async_current_entries(include_ignore=False): + # Try our existing credentials to check for ip change + if controller := await async_get_controller( + self.hass, + ip_address, + entry.data[CONF_PASSWORD], + entry.data[CONF_PORT], + entry.data.get(CONF_SSL, True), + ): + await self.async_set_unique_id(controller.mac) + self._abort_if_unique_id_configured( + updates={CONF_IP_ADDRESS: ip_address} + ) + + # A new rain machine: We will change out the unique id + # for the mac address once we authenticate, however we want to + # prevent multiple different rain machines on the same network + # from being shown in discovery + await self.async_set_unique_id(ip_address) + self._abort_if_unique_id_configured() + self.discovered_ip_address = ip_address + return await self.async_step_user() + + @callback + def _async_generate_schema(self): + """Generate schema.""" + return vol.Schema( + { + vol.Required(CONF_IP_ADDRESS, default=self.discovered_ip_address): str, + vol.Required(CONF_PASSWORD): str, + vol.Optional(CONF_PORT, default=DEFAULT_PORT): int, + } + ) + async def async_step_user(self, user_input=None): """Handle the start of the config flow.""" - if not user_input: - return self.async_show_form( - step_id="user", data_schema=DATA_SCHEMA, errors={} - ) - - await self.async_set_unique_id(user_input[CONF_IP_ADDRESS]) - self._abort_if_unique_id_configured() - - websession = aiohttp_client.async_get_clientsession(self.hass) - client = Client(session=websession) - - try: - await client.load_local( + errors = {} + if user_input: + self._async_abort_ip_address_configured(user_input[CONF_IP_ADDRESS]) + controller = await async_get_controller( + self.hass, user_input[CONF_IP_ADDRESS], user_input[CONF_PASSWORD], - port=user_input[CONF_PORT], - ssl=user_input.get(CONF_SSL, True), - ) - except RainMachineError: - return self.async_show_form( - step_id="user", - data_schema=DATA_SCHEMA, - errors={CONF_PASSWORD: "invalid_auth"}, + user_input[CONF_PORT], + user_input.get(CONF_SSL, True), ) + if controller: + await self.async_set_unique_id(controller.mac) + self._abort_if_unique_id_configured() - # Unfortunately, RainMachine doesn't provide a way to refresh the - # access token without using the IP address and password, so we have to - # store it: - return self.async_create_entry( - title=user_input[CONF_IP_ADDRESS], - data={ - CONF_IP_ADDRESS: user_input[CONF_IP_ADDRESS], - CONF_PASSWORD: user_input[CONF_PASSWORD], - CONF_PORT: user_input[CONF_PORT], - CONF_SSL: user_input.get(CONF_SSL, True), - CONF_ZONE_RUN_TIME: user_input.get( - CONF_ZONE_RUN_TIME, DEFAULT_ZONE_RUN - ), - }, + # Unfortunately, RainMachine doesn't provide a way to refresh the + # access token without using the IP address and password, so we have to + # store it: + return self.async_create_entry( + title=controller.name, + data={ + CONF_IP_ADDRESS: user_input[CONF_IP_ADDRESS], + CONF_PASSWORD: user_input[CONF_PASSWORD], + CONF_PORT: user_input[CONF_PORT], + CONF_SSL: user_input.get(CONF_SSL, True), + CONF_ZONE_RUN_TIME: user_input.get( + CONF_ZONE_RUN_TIME, DEFAULT_ZONE_RUN + ), + }, + ) + + errors = {CONF_PASSWORD: "invalid_auth"} + + if self.discovered_ip_address: + self.context["title_placeholders"] = {"ip": self.discovered_ip_address} + return self.async_show_form( + step_id="user", data_schema=self._async_generate_schema(), errors=errors ) diff --git a/homeassistant/components/rainmachine/const.py b/homeassistant/components/rainmachine/const.py index 568108e23a6..56c1660a0ba 100644 --- a/homeassistant/components/rainmachine/const.py +++ b/homeassistant/components/rainmachine/const.py @@ -9,7 +9,6 @@ CONF_ZONE_RUN_TIME = "zone_run_time" DATA_CONTROLLER = "controller" DATA_COORDINATOR = "coordinator" -DATA_LISTENER = "listener" DATA_PROGRAMS = "programs" DATA_PROVISION_SETTINGS = "provision.settings" DATA_RESTRICTIONS_CURRENT = "restrictions.current" diff --git a/homeassistant/components/rainmachine/manifest.json b/homeassistant/components/rainmachine/manifest.json index 17429a74d40..b6021d02c39 100644 --- a/homeassistant/components/rainmachine/manifest.json +++ b/homeassistant/components/rainmachine/manifest.json @@ -5,5 +5,14 @@ "documentation": "https://www.home-assistant.io/integrations/rainmachine", "requirements": ["regenmaschine==3.0.0"], "codeowners": ["@bachya"], - "iot_class": "local_polling" + "iot_class": "local_polling", + "homekit": { + "models": ["Touch HD", "SPK5"] + }, + "zeroconf": [ + { + "type": "_http._tcp.local.", + "name": "rainmachine*" + } + ] } diff --git a/homeassistant/components/rainmachine/strings.json b/homeassistant/components/rainmachine/strings.json index 1f5a21d37d8..ec65d1c7c09 100644 --- a/homeassistant/components/rainmachine/strings.json +++ b/homeassistant/components/rainmachine/strings.json @@ -1,5 +1,6 @@ { "config": { + "flow_title": "RainMachine {ip}", "step": { "user": { "title": "Fill in your information", diff --git a/homeassistant/components/rainmachine/translations/en.json b/homeassistant/components/rainmachine/translations/en.json index f65463626e4..0c8b6eef766 100644 --- a/homeassistant/components/rainmachine/translations/en.json +++ b/homeassistant/components/rainmachine/translations/en.json @@ -6,6 +6,7 @@ "error": { "invalid_auth": "Invalid authentication" }, + "flow_title": "RainMachine {ip}", "step": { "user": { "data": { diff --git a/homeassistant/generated/zeroconf.py b/homeassistant/generated/zeroconf.py index 33c08579c36..0b1c0adb9c6 100644 --- a/homeassistant/generated/zeroconf.py +++ b/homeassistant/generated/zeroconf.py @@ -98,6 +98,10 @@ ZEROCONF = { "domain": "rachio", "name": "rachio*" }, + { + "domain": "rainmachine", + "name": "rainmachine*" + }, { "domain": "shelly", "name": "shelly*" @@ -217,9 +221,11 @@ HOMEKIT = { "PowerView": "hunterdouglas_powerview", "Presence": "netatmo", "Rachio": "rachio", + "SPK5": "rainmachine", "Smart Bridge": "lutron_caseta", "Socket": "wemo", "TRADFRI": "tradfri", + "Touch HD": "rainmachine", "Welcome": "netatmo", "Wemo": "wemo", "iSmartGate": "gogogate2", diff --git a/tests/components/rainmachine/test_config_flow.py b/tests/components/rainmachine/test_config_flow.py index e79874831fe..1a015a5b181 100644 --- a/tests/components/rainmachine/test_config_flow.py +++ b/tests/components/rainmachine/test_config_flow.py @@ -1,16 +1,25 @@ """Define tests for the OpenUV config flow.""" -from unittest.mock import patch +from unittest.mock import AsyncMock, Mock, patch +import pytest from regenmaschine.errors import RainMachineError -from homeassistant import data_entry_flow -from homeassistant.components.rainmachine import CONF_ZONE_RUN_TIME, DOMAIN, config_flow -from homeassistant.config_entries import SOURCE_USER +from homeassistant import config_entries, data_entry_flow +from homeassistant.components.rainmachine import CONF_ZONE_RUN_TIME, DOMAIN from homeassistant.const import CONF_IP_ADDRESS, CONF_PASSWORD, CONF_PORT, CONF_SSL from tests.common import MockConfigEntry +def _get_mock_client(): + mock_controller = Mock() + mock_controller.name = "My Rain Machine" + mock_controller.mac = "aa:bb:cc:dd:ee:ff" + return Mock( + load_local=AsyncMock(), controllers={"aa:bb:cc:dd:ee:ff": mock_controller} + ) + + async def test_duplicate_error(hass): """Test that errors are shown when duplicates are added.""" conf = { @@ -20,13 +29,19 @@ async def test_duplicate_error(hass): CONF_SSL: True, } - MockConfigEntry(domain=DOMAIN, unique_id="192.168.1.100", data=conf).add_to_hass( - hass - ) + MockConfigEntry( + domain=DOMAIN, unique_id="aa:bb:cc:dd:ee:ff", data=conf + ).add_to_hass(hass) - result = await hass.config_entries.flow.async_init( - DOMAIN, context={"source": SOURCE_USER}, data=conf - ) + with patch( + "homeassistant.components.rainmachine.config_flow.Client", + return_value=_get_mock_client(), + ): + result = await hass.config_entries.flow.async_init( + DOMAIN, + context={"source": config_entries.SOURCE_USER}, + data=conf, + ) assert result["type"] == data_entry_flow.RESULT_TYPE_ABORT assert result["reason"] == "already_configured" @@ -40,16 +55,18 @@ async def test_invalid_password(hass): CONF_SSL: True, } - flow = config_flow.RainMachineFlowHandler() - flow.hass = hass - flow.context = {"source": SOURCE_USER} - with patch( "regenmaschine.client.Client.load_local", side_effect=RainMachineError, ): - result = await flow.async_step_user(user_input=conf) - assert result["errors"] == {CONF_PASSWORD: "invalid_auth"} + result = await hass.config_entries.flow.async_init( + DOMAIN, + context={"source": config_entries.SOURCE_USER}, + data=conf, + ) + await hass.async_block_till_done() + + assert result["errors"] == {CONF_PASSWORD: "invalid_auth"} async def test_options_flow(hass): @@ -88,11 +105,11 @@ async def test_options_flow(hass): async def test_show_form(hass): """Test that the form is served with no input.""" - flow = config_flow.RainMachineFlowHandler() - flow.hass = hass - flow.context = {"source": SOURCE_USER} - - result = await flow.async_step_user(user_input=None) + result = await hass.config_entries.flow.async_init( + DOMAIN, + context={"source": config_entries.SOURCE_USER}, + data=None, + ) assert result["type"] == data_entry_flow.RESULT_TYPE_FORM assert result["step_id"] == "user" @@ -107,22 +124,172 @@ async def test_step_user(hass): CONF_SSL: True, } - flow = config_flow.RainMachineFlowHandler() - flow.hass = hass - flow.context = {"source": SOURCE_USER} + with patch( + "homeassistant.components.rainmachine.async_setup_entry", return_value=True + ) as mock_setup_entry, patch( + "homeassistant.components.rainmachine.config_flow.Client", + return_value=_get_mock_client(), + ): + result = await hass.config_entries.flow.async_init( + DOMAIN, + context={"source": config_entries.SOURCE_USER}, + data=conf, + ) + + assert result["type"] == data_entry_flow.RESULT_TYPE_CREATE_ENTRY + assert result["title"] == "My Rain Machine" + assert result["data"] == { + CONF_IP_ADDRESS: "192.168.1.100", + CONF_PASSWORD: "password", + CONF_PORT: 8080, + CONF_SSL: True, + CONF_ZONE_RUN_TIME: 600, + } + assert mock_setup_entry.called + + +@pytest.mark.parametrize( + "source", [config_entries.SOURCE_ZEROCONF, config_entries.SOURCE_HOMEKIT] +) +async def test_step_homekit_zeroconf_ip_already_exists(hass, source): + """Test homekit and zeroconf with an ip that already exists.""" + conf = { + CONF_IP_ADDRESS: "192.168.1.100", + CONF_PASSWORD: "password", + CONF_PORT: 8080, + CONF_SSL: True, + } + + MockConfigEntry( + domain=DOMAIN, unique_id="aa:bb:cc:dd:ee:ff", data=conf + ).add_to_hass(hass) with patch( - "regenmaschine.client.Client.load_local", - return_value=True, + "homeassistant.components.rainmachine.config_flow.Client", + return_value=_get_mock_client(), ): - result = await flow.async_step_user(user_input=conf) + result = await hass.config_entries.flow.async_init( + DOMAIN, + context={"source": source}, + data={"host": "192.168.1.100"}, + ) - assert result["type"] == data_entry_flow.RESULT_TYPE_CREATE_ENTRY - assert result["title"] == "192.168.1.100" - assert result["data"] == { - CONF_IP_ADDRESS: "192.168.1.100", - CONF_PASSWORD: "password", - CONF_PORT: 8080, - CONF_SSL: True, - CONF_ZONE_RUN_TIME: 600, - } + assert result["type"] == data_entry_flow.RESULT_TYPE_ABORT + assert result["reason"] == "already_configured" + + +@pytest.mark.parametrize( + "source", [config_entries.SOURCE_ZEROCONF, config_entries.SOURCE_HOMEKIT] +) +async def test_step_homekit_zeroconf_ip_change(hass, source): + """Test zeroconf with an ip change.""" + conf = { + CONF_IP_ADDRESS: "192.168.1.100", + CONF_PASSWORD: "password", + CONF_PORT: 8080, + CONF_SSL: True, + } + + entry = MockConfigEntry(domain=DOMAIN, unique_id="aa:bb:cc:dd:ee:ff", data=conf) + entry.add_to_hass(hass) + + with patch( + "homeassistant.components.rainmachine.config_flow.Client", + return_value=_get_mock_client(), + ): + result = await hass.config_entries.flow.async_init( + DOMAIN, + context={"source": source}, + data={"host": "192.168.1.2"}, + ) + + assert result["type"] == data_entry_flow.RESULT_TYPE_ABORT + assert result["reason"] == "already_configured" + assert entry.data[CONF_IP_ADDRESS] == "192.168.1.2" + + +@pytest.mark.parametrize( + "source", [config_entries.SOURCE_ZEROCONF, config_entries.SOURCE_HOMEKIT] +) +async def test_step_homekit_zeroconf_new_controller_when_some_exist(hass, source): + """Test homekit and zeroconf for a new controller when one already exists.""" + existing_conf = { + CONF_IP_ADDRESS: "192.168.1.3", + CONF_PASSWORD: "password", + CONF_PORT: 8080, + CONF_SSL: True, + } + entry = MockConfigEntry( + domain=DOMAIN, unique_id="zz:bb:cc:dd:ee:ff", data=existing_conf + ) + entry.add_to_hass(hass) + + with patch( + "homeassistant.components.rainmachine.config_flow.Client", + return_value=_get_mock_client(), + ): + result = await hass.config_entries.flow.async_init( + DOMAIN, + context={"source": source}, + data={"host": "192.168.1.100"}, + ) + + assert result["type"] == data_entry_flow.RESULT_TYPE_FORM + assert result["step_id"] == "user" + + with patch( + "homeassistant.components.rainmachine.async_setup_entry", return_value=True + ) as mock_setup_entry, patch( + "homeassistant.components.rainmachine.config_flow.Client", + return_value=_get_mock_client(), + ): + result2 = await hass.config_entries.flow.async_configure( + result["flow_id"], + { + CONF_IP_ADDRESS: "192.168.1.100", + CONF_PASSWORD: "password", + CONF_PORT: 8080, + }, + ) + await hass.async_block_till_done() + + assert result2["type"] == data_entry_flow.RESULT_TYPE_CREATE_ENTRY + assert result2["title"] == "My Rain Machine" + assert result2["data"] == { + CONF_IP_ADDRESS: "192.168.1.100", + CONF_PASSWORD: "password", + CONF_PORT: 8080, + CONF_SSL: True, + CONF_ZONE_RUN_TIME: 600, + } + assert mock_setup_entry.called + + +async def test_discovery_by_homekit_and_zeroconf_same_time(hass): + """Test the same controller gets discovered by two different methods.""" + + with patch( + "homeassistant.components.rainmachine.config_flow.Client", + return_value=_get_mock_client(), + ): + result = await hass.config_entries.flow.async_init( + DOMAIN, + context={"source": config_entries.SOURCE_ZEROCONF}, + data={"host": "192.168.1.100"}, + ) + + assert result["type"] == data_entry_flow.RESULT_TYPE_FORM + assert result["step_id"] == "user" + + with patch( + "homeassistant.components.rainmachine.config_flow.Client", + return_value=_get_mock_client(), + ): + result2 = await hass.config_entries.flow.async_init( + DOMAIN, + context={"source": config_entries.SOURCE_HOMEKIT}, + data={"host": "192.168.1.100"}, + ) + + assert result2["type"] == data_entry_flow.RESULT_TYPE_ABORT + assert result2["reason"] == "already_in_progress"