Add model information for probe_plus devices (home-assistant#154262)

Author: pantherale0

homeassistant/components/probe_plus/__init__.py

@@ -2,7 +2,7 @@
 
 from __future__ import annotations
 
-from homeassistant.const import Platform
+from homeassistant.const import CONF_MODEL, Platform
 from homeassistant.core import HomeAssistant
 
 from .coordinator import ProbePlusConfigEntry, ProbePlusDataUpdateCoordinator
@@ -12,6 +12,12 @@
 
 async def async_setup_entry(hass: HomeAssistant, entry: ProbePlusConfigEntry) -> bool:
     """Set up Probe Plus from a config entry."""
+    # Perform a migration to ensure the model is added to the config entry schema.
+    if CONF_MODEL not in entry.data:
+        # The config entry adds the model number of the device to the start of its title
+        hass.config_entries.async_update_entry(
+            entry, data={**entry.data, CONF_MODEL: entry.title.split(" ")[0]}
+        )
     coordinator = ProbePlusDataUpdateCoordinator(hass, entry)
     await coordinator.async_config_entry_first_refresh()
     entry.runtime_data = coordinator

homeassistant/components/probe_plus/config_flow.py

@@ -13,7 +13,7 @@
     async_discovered_service_info,
 )
 from homeassistant.config_entries import ConfigFlow, ConfigFlowResult
-from homeassistant.const import CONF_ADDRESS
+from homeassistant.const import CONF_ADDRESS, CONF_MODEL
 
 from .const import DOMAIN
 
@@ -73,6 +73,7 @@ async def async_step_bluetooth_confirm(
                 title=discovery.title,
                 data={
                     CONF_ADDRESS: discovery.discovery_info.address,
+                    CONF_MODEL: discovery.discovery_info.name,
                 },
             )
         self._set_confirm_only()
@@ -95,7 +96,7 @@ async def async_step_user(
             discovery = self._discovered_devices[address]
             return self.async_create_entry(
                 title=discovery.title,
-                data=user_input,
+                data={**user_input, CONF_MODEL: discovery.discovery_info.name},
             )
 
         current_addresses = self._async_current_ids()

homeassistant/components/probe_plus/entity.py

@@ -4,6 +4,7 @@
 
 from pyprobeplus import ProbePlusDevice
 
+from homeassistant.const import CONF_MODEL
 from homeassistant.helpers.device_registry import (
     CONNECTION_BLUETOOTH,
     DeviceInfo,
@@ -40,6 +41,7 @@ def __init__(
             name=coordinator.device.name,
             manufacturer="Probe Plus",
             suggested_area="Kitchen",
+            model=coordinator.config_entry.data.get(CONF_MODEL),
             connections={(CONNECTION_BLUETOOTH, coordinator.device.mac)},
         )
 

tests/components/probe_plus/test_config_flow.py

@@ -7,7 +7,7 @@
 
 from homeassistant.components.probe_plus.const import DOMAIN
 from homeassistant.config_entries import SOURCE_BLUETOOTH, SOURCE_USER
-from homeassistant.const import CONF_ADDRESS
+from homeassistant.const import CONF_ADDRESS, CONF_MODEL
 from homeassistant.core import HomeAssistant
 from homeassistant.data_entry_flow import FlowResultType
 from homeassistant.helpers.service_info.bluetooth import BluetoothServiceInfo
@@ -57,7 +57,7 @@ async def test_user_config_flow_creates_entry(
     assert result["type"] is FlowResultType.CREATE_ENTRY
     assert result["result"].unique_id == "aa:bb:cc:dd:ee:ff"
     assert result["title"] == "FM210 aa:bb:cc:dd:ee:ff"
-    assert result["data"] == {CONF_ADDRESS: "aa:bb:cc:dd:ee:ff"}
+    assert result["data"] == {CONF_ADDRESS: "aa:bb:cc:dd:ee:ff", CONF_MODEL: "FM210"}
 
 
 async def test_user_flow_already_configured(
@@ -97,9 +97,7 @@ async def test_bluetooth_discovery(
     assert result["type"] is FlowResultType.CREATE_ENTRY
     assert result["title"] == "FM210 aa:bb:cc:dd:ee:ff"
     assert result["result"].unique_id == "aa:bb:cc:dd:ee:ff"
-    assert result["data"] == {
-        CONF_ADDRESS: service_info.address,
-    }
+    assert result["data"] == {CONF_ADDRESS: service_info.address, CONF_MODEL: "FM210"}
 
 
 async def test_already_configured_bluetooth_discovery(