Add rainmachine discovery (#49970)

Co-authored-by: Paulus Schoutsen <paulus@home-assistant.io>

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(
+    controller = hass.data[DOMAIN][DATA_CONTROLLER][
-        iter(client.controllers.values())
+        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
 
 

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,49 +21,112 @@ 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:
+        errors = {}
-            return self.async_show_form(
+        if user_input:
-                step_id="user", data_schema=DATA_SCHEMA, errors={}
+            self._async_abort_ip_address_configured(user_input[CONF_IP_ADDRESS])
-            )
+            controller = await async_get_controller(
-
+                self.hass,
-        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(
                 user_input[CONF_IP_ADDRESS],
                 user_input[CONF_PASSWORD],
-                port=user_input[CONF_PORT],
+                user_input[CONF_PORT],
-                ssl=user_input.get(CONF_SSL, True),
+                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"},
             )
+            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],
+                    title=controller.name,
                     data={
                         CONF_IP_ADDRESS: user_input[CONF_IP_ADDRESS],
                         CONF_PASSWORD: user_input[CONF_PASSWORD],
@@ -73,6 +138,14 @@ class RainMachineFlowHandler(config_entries.ConfigFlow, domain=DOMAIN):
                     },
                 )
 
+            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
+        )
+
 
 class RainMachineOptionsFlowHandler(config_entries.OptionsFlow):
     """Handle a RainMachine options flow."""

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"

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*"
+    }
+  ]
 }

homeassistant/components/rainmachine/strings.json

@@ -1,5 +1,6 @@
 {
   "config": {
+    "flow_title": "RainMachine {ip}",
     "step": {
       "user": {
         "title": "Fill in your information",

homeassistant/components/rainmachine/translations/en.json

@@ -6,6 +6,7 @@
         "error": {
             "invalid_auth": "Invalid authentication"
         },
+        "flow_title": "RainMachine {ip}",
         "step": {
             "user": {
                 "data": {

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",

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 import config_entries, data_entry_flow
-from homeassistant.components.rainmachine import CONF_ZONE_RUN_TIME, DOMAIN, config_flow
+from homeassistant.components.rainmachine import CONF_ZONE_RUN_TIME, DOMAIN
-from homeassistant.config_entries import SOURCE_USER
 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,12 +29,18 @@ async def test_duplicate_error(hass):
         CONF_SSL: True,
     }
 
-    MockConfigEntry(domain=DOMAIN, unique_id="192.168.1.100", data=conf).add_to_hass(
+    MockConfigEntry(
-        hass
+        domain=DOMAIN, unique_id="aa:bb:cc:dd:ee:ff", data=conf
-    )
+    ).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_USER}, data=conf
+            DOMAIN,
+            context={"source": config_entries.SOURCE_USER},
+            data=conf,
         )
     assert result["type"] == data_entry_flow.RESULT_TYPE_ABORT
     assert result["reason"] == "already_configured"
@@ -40,15 +55,17 @@ 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)
+        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"}
 
 
@@ -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()
+    result = await hass.config_entries.flow.async_init(
-    flow.hass = hass
+        DOMAIN,
-    flow.context = {"source": SOURCE_USER}
+        context={"source": config_entries.SOURCE_USER},
-
+        data=None,
-    result = await flow.async_step_user(user_input=None)
+    )
 
     assert result["type"] == data_entry_flow.RESULT_TYPE_FORM
     assert result["step_id"] == "user"
@@ -107,18 +124,20 @@ async def test_step_user(hass):
         CONF_SSL: True,
     }
 
-    flow = config_flow.RainMachineFlowHandler()
-    flow.hass = hass
-    flow.context = {"source": SOURCE_USER}
-
     with patch(
-        "regenmaschine.client.Client.load_local",
+        "homeassistant.components.rainmachine.async_setup_entry", return_value=True
-        return_value=True,
+    ) as mock_setup_entry, patch(
+        "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": config_entries.SOURCE_USER},
+            data=conf,
+        )
 
     assert result["type"] == data_entry_flow.RESULT_TYPE_CREATE_ENTRY
-        assert result["title"] == "192.168.1.100"
+    assert result["title"] == "My Rain Machine"
     assert result["data"] == {
         CONF_IP_ADDRESS: "192.168.1.100",
         CONF_PASSWORD: "password",
@@ -126,3 +145,151 @@ async def test_step_user(hass):
         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(
+        "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_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"