Add support for discovering individual roombas (#45200)

* Add support for discovering individual roombas

* add missing translation string

* Update homeassistant/components/roomba/strings.json

Co-authored-by: Martin Hjelmare <marhje52@gmail.com>

Co-authored-by: Martin Hjelmare <marhje52@gmail.com>

homeassistant/components/roomba/config_flow.py

@@ -1,11 +1,14 @@
 """Config flow to configure roomba component."""
 
+import asyncio
+
 from roombapy import Roomba
 from roombapy.discovery import RoombaDiscovery
 from roombapy.getpassword import RoombaPassword
 import voluptuous as vol
 
 from homeassistant import config_entries, core
+from homeassistant.components.dhcp import HOSTNAME, IP_ADDRESS
 from homeassistant.const import CONF_HOST, CONF_PASSWORD
 from homeassistant.core import callback
 
@@ -21,6 +24,8 @@ from .const import (
 )
 from .const import DOMAIN  # pylint:disable=unused-import
 
+ROOMBA_DISCOVERY_LOCK = "roomba_discovery_lock"
+
 DEFAULT_OPTIONS = {CONF_CONTINUOUS: DEFAULT_CONTINUOUS, CONF_DELAY: DEFAULT_DELAY}
 
 MAX_NUM_DEVICES_TO_DISCOVER = 25
@@ -72,6 +77,35 @@ class RoombaConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
         """Get the options flow for this handler."""
         return OptionsFlowHandler(config_entry)
 
+    async def async_step_dhcp(self, dhcp_discovery):
+        """Handle dhcp discovery."""
+        if self._async_host_already_configured(dhcp_discovery[IP_ADDRESS]):
+            return self.async_abort(reason="already_configured")
+
+        if not dhcp_discovery[HOSTNAME].startswith("iRobot-"):
+            return self.async_abort(reason="not_irobot_device")
+
+        blid = _async_blid_from_hostname(dhcp_discovery[HOSTNAME])
+        await self.async_set_unique_id(blid)
+        self._abort_if_unique_id_configured(
+            updates={CONF_HOST: dhcp_discovery[IP_ADDRESS]}
+        )
+
+        self.host = dhcp_discovery[IP_ADDRESS]
+        self.blid = blid
+        # pylint: disable=no-member # https://github.com/PyCQA/pylint/issues/3167
+        self.context["title_placeholders"] = {"host": self.host, "name": self.blid}
+        return await self.async_step_user()
+
+    async def _async_start_link(self):
+        """Start linking."""
+        device = self.discovered_robots[self.host]
+        self.blid = device.blid
+        self.name = device.robot_name
+        await self.async_set_unique_id(self.blid, raise_on_progress=False)
+        self._abort_if_unique_id_configured()
+        return await self.async_step_link()
+
     async def async_step_user(self, user_input=None):
         """Handle a flow start."""
         # Check if user chooses manual entry
@@ -84,15 +118,12 @@ class RoombaConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
             and user_input[CONF_HOST] in self.discovered_robots
         ):
             self.host = user_input[CONF_HOST]
-            device = self.discovered_robots[self.host]
+            return await self._async_start_link()
-            self.blid = device.blid
-            self.name = device.robot_name
-            await self.async_set_unique_id(self.blid, raise_on_progress=False)
-            self._abort_if_unique_id_configured()
-            return await self.async_step_link()
 
         already_configured = self._async_current_ids(False)
         discovery = _async_get_roomba_discovery()
+
+        async with self.hass.data.setdefault(ROOMBA_DISCOVERY_LOCK, asyncio.Lock()):
             devices = await self.hass.async_add_executor_job(discovery.get_all)
 
         if devices:
@@ -102,6 +133,14 @@ class RoombaConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
                 for device in devices
                 if device.blid not in already_configured
             }
+            if self.host and self.host in self.discovered_robots:
+                # From discovery
+                # pylint: disable=no-member # https://github.com/PyCQA/pylint/issues/3167
+                self.context["title_placeholders"] = {
+                    "host": self.host,
+                    "name": self.discovered_robots[self.host].robot_name,
+                }
+                return await self._async_start_link()
 
         if not self.discovered_robots:
             return await self.async_step_manual()
@@ -131,7 +170,10 @@ class RoombaConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
                 step_id="manual",
                 description_placeholders={AUTH_HELP_URL_KEY: AUTH_HELP_URL_VALUE},
                 data_schema=vol.Schema(
-                    {vol.Required(CONF_HOST): str, vol.Required(CONF_BLID): str}
+                    {
+                        vol.Required(CONF_HOST, default=self.host): str,
+                        vol.Required(CONF_BLID, default=self.blid): str,
+                    }
                 ),
             )
 
@@ -154,7 +196,10 @@ class RoombaConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
         to connect to the device.
         """
         if user_input is None:
-            return self.async_show_form(step_id="link")
+            return self.async_show_form(
+                step_id="link",
+                description_placeholders={CONF_NAME: self.name or self.blid},
+            )
 
         try:
             password = await self.hass.async_add_executor_job(
@@ -211,6 +256,14 @@ class RoombaConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
             errors=errors,
         )
 
+    @callback
+    def _async_host_already_configured(self, host):
+        """See if we already have an entry matching the host."""
+        for entry in self._async_current_entries():
+            if entry.data.get(CONF_HOST) == host:
+                return True
+        return False
+
 
 class OptionsFlowHandler(config_entries.OptionsFlow):
     """Handle options."""
@@ -251,3 +304,9 @@ def _async_get_roomba_discovery():
     discovery = RoombaDiscovery()
     discovery.amount_of_broadcasted_messages = MAX_NUM_DEVICES_TO_DISCOVER
     return discovery
+
+
+@callback
+def _async_blid_from_hostname(hostname):
+    """Extract the blid from the hostname."""
+    return hostname.split("-")[1].split(".")[0]

homeassistant/components/roomba/strings.json

@@ -1,5 +1,6 @@
 {
   "config": {
+    "flow_title": "iRobot {name} ({host})",
     "step": {
       "init": {
         "title": "Automaticlly connect to the device",
@@ -18,7 +19,7 @@
       },      
       "link": {
         "title": "Retrieve Password",
-        "description": "Press and hold the Home button until the device generates a sound (about two seconds)."
+        "description": "Press and hold the Home button on {name} until the device generates a sound (about two seconds)."
       },
       "link_manual": {
         "title": "Enter Password",
@@ -32,7 +33,9 @@
       "cannot_connect": "[%key:common::config_flow::error::cannot_connect%]"
     },
     "abort": {
-      "cannot_connect": "[%key:common::config_flow::error::cannot_connect%]"
+      "cannot_connect": "[%key:common::config_flow::error::cannot_connect%]",
+      "already_configured": "[%key:common::config_flow::abort::already_configured_device%]",
+      "not_irobot_device": "Discovered device is not an iRobot device"
     }    
   },
   "options": {

homeassistant/components/roomba/translations/en.json

@@ -1,5 +1,6 @@
 {
   "config": {
+    "flow_title": "iRobot {name} ({host})",
     "step": {
       "init": {
         "title": "Automaticlly connect to the device",
@@ -18,7 +19,7 @@
       },      
       "link": {
         "title": "Retrieve Password",
-        "description": "Press and hold the Home button until the device generates a sound (about two seconds)."
+        "description": "Press and hold the Home button on {name} until the device generates a sound (about two seconds)."
       },
       "link_manual": {
         "title": "Enter Password",
@@ -32,7 +33,9 @@
       "cannot_connect": "[%key:common::config_flow::error::cannot_connect%]"
     },
     "abort": {
-      "cannot_connect": "[%key:common::config_flow::error::cannot_connect%]"
+      "cannot_connect": "[%key:common::config_flow::error::cannot_connect%]",
+      "already_configured": "[%key:common::config_flow::abort::already_configured_device%]",
+      "not_irobot_device": "Discovered device not an iRobot device"
     }    
   },
   "options": {

tests/components/roomba/test_config_flow.py

@@ -5,6 +5,7 @@ from roombapy import RoombaConnectionError
 from roombapy.roomba import RoombaInfo
 
 from homeassistant import config_entries, data_entry_flow, setup
+from homeassistant.components.dhcp import HOSTNAME, IP_ADDRESS, MAC_ADDRESS
 from homeassistant.components.roomba.const import (
     CONF_BLID,
     CONF_CONTINUOUS,
@@ -579,3 +580,242 @@ async def test_form_user_discovery_and_password_fetch_gets_connection_refused(ha
     }
     assert len(mock_setup.mock_calls) == 1
     assert len(mock_setup_entry.mock_calls) == 1
+
+
+async def test_dhcp_discovery_and_roomba_discovery_finds(hass):
+    """Test we can process the discovery from dhcp and roomba discovery matches the device."""
+    await setup.async_setup_component(hass, "persistent_notification", {})
+
+    mocked_roomba = _create_mocked_roomba(
+        roomba_connected=True,
+        master_state={"state": {"reported": {"name": "myroomba"}}},
+    )
+
+    with patch(
+        "homeassistant.components.roomba.config_flow.RoombaDiscovery", _mocked_discovery
+    ):
+        result = await hass.config_entries.flow.async_init(
+            DOMAIN,
+            context={"source": config_entries.SOURCE_DHCP},
+            data={
+                IP_ADDRESS: MOCK_IP,
+                MAC_ADDRESS: "AA:BB:CC:DD:EE:FF",
+                HOSTNAME: "iRobot-blid",
+            },
+        )
+        await hass.async_block_till_done()
+
+    assert result["type"] == data_entry_flow.RESULT_TYPE_FORM
+    assert result["errors"] is None
+    assert result["step_id"] == "link"
+    assert result["description_placeholders"] == {"name": "robot_name"}
+
+    with patch(
+        "homeassistant.components.roomba.config_flow.Roomba",
+        return_value=mocked_roomba,
+    ), patch(
+        "homeassistant.components.roomba.config_flow.RoombaPassword",
+        _mocked_getpassword,
+    ), patch(
+        "homeassistant.components.roomba.async_setup", return_value=True
+    ) as mock_setup, patch(
+        "homeassistant.components.roomba.async_setup_entry",
+        return_value=True,
+    ) as mock_setup_entry:
+        result2 = await hass.config_entries.flow.async_configure(
+            result["flow_id"],
+            {},
+        )
+        await hass.async_block_till_done()
+
+    assert result2["type"] == data_entry_flow.RESULT_TYPE_CREATE_ENTRY
+    assert result2["title"] == "robot_name"
+    assert result2["result"].unique_id == "blid"
+    assert result2["data"] == {
+        CONF_BLID: "blid",
+        CONF_CONTINUOUS: True,
+        CONF_DELAY: 1,
+        CONF_HOST: MOCK_IP,
+        CONF_PASSWORD: "password",
+    }
+    assert len(mock_setup.mock_calls) == 1
+    assert len(mock_setup_entry.mock_calls) == 1
+
+
+async def test_dhcp_discovery_falls_back_to_manual(hass):
+    """Test we can process the discovery from dhcp but roomba discovery cannot find the device."""
+    await setup.async_setup_component(hass, "persistent_notification", {})
+
+    mocked_roomba = _create_mocked_roomba(
+        roomba_connected=True,
+        master_state={"state": {"reported": {"name": "myroomba"}}},
+    )
+
+    with patch(
+        "homeassistant.components.roomba.config_flow.RoombaDiscovery", _mocked_discovery
+    ):
+        result = await hass.config_entries.flow.async_init(
+            DOMAIN,
+            context={"source": config_entries.SOURCE_DHCP},
+            data={
+                IP_ADDRESS: "1.1.1.1",
+                MAC_ADDRESS: "AA:BB:CC:DD:EE:FF",
+                HOSTNAME: "iRobot-blid",
+            },
+        )
+        await hass.async_block_till_done()
+
+    assert result["type"] == data_entry_flow.RESULT_TYPE_FORM
+    assert result["errors"] is None
+    assert result["step_id"] == "user"
+
+    result2 = await hass.config_entries.flow.async_configure(
+        result["flow_id"],
+        {},
+    )
+    await hass.async_block_till_done()
+    assert result2["type"] == data_entry_flow.RESULT_TYPE_FORM
+    assert result2["errors"] is None
+    assert result2["step_id"] == "manual"
+
+    result3 = await hass.config_entries.flow.async_configure(
+        result2["flow_id"],
+        {CONF_HOST: "1.1.1.1", CONF_BLID: "blid"},
+    )
+    await hass.async_block_till_done()
+    assert result3["type"] == data_entry_flow.RESULT_TYPE_FORM
+    assert result3["errors"] is None
+
+    with patch(
+        "homeassistant.components.roomba.config_flow.Roomba",
+        return_value=mocked_roomba,
+    ), patch(
+        "homeassistant.components.roomba.config_flow.RoombaPassword",
+        _mocked_getpassword,
+    ), patch(
+        "homeassistant.components.roomba.async_setup", return_value=True
+    ) as mock_setup, patch(
+        "homeassistant.components.roomba.async_setup_entry",
+        return_value=True,
+    ) as mock_setup_entry:
+        result4 = await hass.config_entries.flow.async_configure(
+            result3["flow_id"],
+            {},
+        )
+        await hass.async_block_till_done()
+
+    assert result4["type"] == data_entry_flow.RESULT_TYPE_CREATE_ENTRY
+    assert result4["title"] == "myroomba"
+    assert result4["result"].unique_id == "blid"
+    assert result4["data"] == {
+        CONF_BLID: "blid",
+        CONF_CONTINUOUS: True,
+        CONF_DELAY: 1,
+        CONF_HOST: "1.1.1.1",
+        CONF_PASSWORD: "password",
+    }
+    assert len(mock_setup.mock_calls) == 1
+    assert len(mock_setup_entry.mock_calls) == 1
+
+
+async def test_dhcp_discovery_with_ignored(hass):
+    """Test ignored entries do not break checking for existing entries."""
+    await setup.async_setup_component(hass, "persistent_notification", {})
+
+    config_entry = MockConfigEntry(domain=DOMAIN, data={}, source="ignore")
+    config_entry.add_to_hass(hass)
+
+    with patch(
+        "homeassistant.components.roomba.config_flow.RoombaDiscovery", _mocked_discovery
+    ):
+        result = await hass.config_entries.flow.async_init(
+            DOMAIN,
+            context={"source": config_entries.SOURCE_DHCP},
+            data={
+                IP_ADDRESS: "1.1.1.1",
+                MAC_ADDRESS: "AA:BB:CC:DD:EE:FF",
+                HOSTNAME: "iRobot-blid",
+            },
+        )
+        await hass.async_block_till_done()
+
+    assert result["type"] == "form"
+
+
+async def test_dhcp_discovery_already_configured_host(hass):
+    """Test we abort if the host is already configured."""
+    await setup.async_setup_component(hass, "persistent_notification", {})
+
+    config_entry = MockConfigEntry(domain=DOMAIN, data={CONF_HOST: "1.1.1.1"})
+    config_entry.add_to_hass(hass)
+
+    with patch(
+        "homeassistant.components.roomba.config_flow.RoombaDiscovery", _mocked_discovery
+    ):
+        result = await hass.config_entries.flow.async_init(
+            DOMAIN,
+            context={"source": config_entries.SOURCE_DHCP},
+            data={
+                IP_ADDRESS: "1.1.1.1",
+                MAC_ADDRESS: "AA:BB:CC:DD:EE:FF",
+                HOSTNAME: "iRobot-blid",
+            },
+        )
+        await hass.async_block_till_done()
+
+    assert result["type"] == "abort"
+    assert result["reason"] == "already_configured"
+
+
+async def test_dhcp_discovery_already_configured_blid(hass):
+    """Test we abort if the blid is already configured."""
+    await setup.async_setup_component(hass, "persistent_notification", {})
+
+    config_entry = MockConfigEntry(
+        domain=DOMAIN, data={CONF_BLID: "blid"}, unique_id="blid"
+    )
+    config_entry.add_to_hass(hass)
+
+    with patch(
+        "homeassistant.components.roomba.config_flow.RoombaDiscovery", _mocked_discovery
+    ):
+        result = await hass.config_entries.flow.async_init(
+            DOMAIN,
+            context={"source": config_entries.SOURCE_DHCP},
+            data={
+                IP_ADDRESS: "1.1.1.1",
+                MAC_ADDRESS: "AA:BB:CC:DD:EE:FF",
+                HOSTNAME: "iRobot-blid",
+            },
+        )
+        await hass.async_block_till_done()
+
+    assert result["type"] == "abort"
+    assert result["reason"] == "already_configured"
+
+
+async def test_dhcp_discovery_not_irobot(hass):
+    """Test we abort if the discovered device is not an irobot device."""
+    await setup.async_setup_component(hass, "persistent_notification", {})
+
+    config_entry = MockConfigEntry(
+        domain=DOMAIN, data={CONF_BLID: "blid"}, unique_id="blid"
+    )
+    config_entry.add_to_hass(hass)
+
+    with patch(
+        "homeassistant.components.roomba.config_flow.RoombaDiscovery", _mocked_discovery
+    ):
+        result = await hass.config_entries.flow.async_init(
+            DOMAIN,
+            context={"source": config_entries.SOURCE_DHCP},
+            data={
+                IP_ADDRESS: "1.1.1.1",
+                MAC_ADDRESS: "AA:BB:CC:DD:EE:FF",
+                HOSTNAME: "NotiRobot-blid",
+            },
+        )
+        await hass.async_block_till_done()
+
+    assert result["type"] == "abort"
+    assert result["reason"] == "not_irobot_device"