Fixes for discovery flow

custom_components/powercalc/__init__.py

@@ -383,6 +383,8 @@ def _init_entity_discovery(
                 CONF_ENTITY_ID: source_entity.entity_id,
                 CONF_MANUFACTURER: power_profile.manufacturer,
                 CONF_MODEL: power_profile.model,
+                DISCOVERY_SOURCE_ENTITY: source_entity,
+                DISCOVERY_POWER_PROFILE: power_profile,
             },
         )
 

custom_components/powercalc/config_flow.py

@@ -60,6 +60,8 @@
     CONF_VALUE,
     CONF_VALUE_TEMPLATE,
     CONF_WLED,
+    DISCOVERY_SOURCE_ENTITY,
+    DISCOVERY_POWER_PROFILE,
     DOMAIN,
     ENERGY_INTEGRATION_METHOD_LEFT,
     ENERGY_INTEGRATION_METHODS,
@@ -220,17 +222,22 @@ async def async_step_integration_discovery(
         self.skip_advanced_step = (
             True  # We don't want to ask advanced option when discovered
         )
-        self.sensor_config.update(discovery_info)
+
         self.selected_sensor_type = SensorType.VIRTUAL_POWER
         self.name = discovery_info[CONF_NAME]
         unique_id = discovery_info[CONF_UNIQUE_ID]
         await self.async_set_unique_id(unique_id)
         self._abort_if_unique_id_configured()
 
+        sensor_config = discovery_info.copy()
+
         self.source_entity_id = discovery_info[CONF_ENTITY_ID]
-        self.source_entity = await create_source_entity(
-            self.source_entity_id, self.hass
-        )
+        self.source_entity = discovery_info[DISCOVERY_SOURCE_ENTITY]
+        self.power_profile = discovery_info[DISCOVERY_POWER_PROFILE]
+
+        del sensor_config[DISCOVERY_SOURCE_ENTITY]
+        del sensor_config[DISCOVERY_POWER_PROFILE]
+        self.sensor_config.update(sensor_config)
 
         self.context["title_placeholders"] = {
             "name": self.sensor_config.get(CONF_NAME),
@@ -379,6 +386,7 @@ async def async_step_library(self, user_input: dict[str, str] = None) -> FlowRes
         Ask the user to confirm this or forward to manual library selection
         """
         if user_input is not None:
+            _LOGGER.debug("discovery flow 1")
             if user_input.get(CONF_CONFIRM_AUTODISCOVERED_MODEL) and self.power_profile:
                 self.sensor_config.update(
                     {
@@ -390,12 +398,14 @@ async def async_step_library(self, user_input: dict[str, str] = None) -> FlowRes
 
             return await self.async_step_manufacturer()
 
-        if self.source_entity.entity_entry:
+        if self.source_entity.entity_entry and self.power_profile is None:
             try:
                 self.power_profile = await get_power_profile(
                     self.hass, {}, self.source_entity.entity_entry
                 )
+                _LOGGER.debug("discovery flow 2")
             except ModelNotSupported:
+                _LOGGER.debug("discovery flow 3")
                 self.power_profile = None
         if self.power_profile:
             return self.async_show_form(
@@ -407,7 +417,7 @@ async def async_step_library(self, user_input: dict[str, str] = None) -> FlowRes
                 data_schema=SCHEMA_POWER_AUTODISCOVERED,
                 errors={},
             )
-
+        _LOGGER.debug("discovery flow 4")
         return await self.async_step_manufacturer()
 
     async def async_step_manufacturer(

tests/test_config_flow.py

@@ -51,12 +51,16 @@
     CONF_VALUE,
     CONF_VOLTAGE,
     CONF_WLED,
+    DISCOVERY_SOURCE_ENTITY,
+    DISCOVERY_POWER_PROFILE,
     ENERGY_INTEGRATION_METHOD_LEFT,
     CalculationStrategy,
     SensorType,
 )
 from custom_components.powercalc.errors import StrategyConfigurationError
 from custom_components.test.light import MockLight
+from custom_components.powercalc.common import create_source_entity
+from custom_components.powercalc.power_profile.model_discovery import get_power_profile
 
 from .common import (
     MockConfigEntry,
@@ -74,6 +78,9 @@ async def test_discovery_flow(hass: HomeAssistant):
     light_entity.model = "LCT010"
     await create_mock_light_entity(hass, light_entity)
 
+    source_entity = await create_source_entity(DEFAULT_ENTITY_ID, hass)
+    power_profile = await get_power_profile(hass, {}, source_entity.entity_entry)
+
     result: FlowResult = await hass.config_entries.flow.async_init(
         DOMAIN,
         context={"source": config_entries.SOURCE_INTEGRATION_DISCOVERY},
@@ -83,6 +90,8 @@ async def test_discovery_flow(hass: HomeAssistant):
             CONF_ENTITY_ID: DEFAULT_ENTITY_ID,
             CONF_MANUFACTURER: "signify",
             CONF_MODEL: "LCT010",
+            DISCOVERY_SOURCE_ENTITY: source_entity,
+            DISCOVERY_POWER_PROFILE: power_profile
         },
     )
 

tests/test_init.py

@@ -51,7 +51,7 @@
 def mock_flow_init(hass):
     """Mock hass.config_entries.flow.async_init."""
     with patch.object(
-        hass.config_entries.flow, "async_init", return_value=AsyncMock()
+            hass.config_entries.flow, "async_init", return_value=AsyncMock()
     ) as mock_init:
         yield mock_init
 
@@ -77,30 +77,12 @@ async def test_autodiscovery(hass: HomeAssistant, mock_flow_init):
 
     # Check that two discovery flows have been initialized
     # LightA and LightB should be discovered, LightC not
-    assert mock_flow_init.mock_calls == [
-        call(
-            DOMAIN,
-            context={"source": SOURCE_INTEGRATION_DISCOVERY},
-            data={
-                CONF_UNIQUE_ID: lighta.unique_id,
-                CONF_NAME: lighta.name,
-                CONF_ENTITY_ID: lighta.entity_id,
-                CONF_MANUFACTURER: lighta.manufacturer,
-                CONF_MODEL: lighta.model,
-            },
-        ),
-        call(
-            DOMAIN,
-            context={"source": SOURCE_INTEGRATION_DISCOVERY},
-            data={
-                CONF_UNIQUE_ID: lightb.unique_id,
-                CONF_NAME: lightb.name,
-                CONF_ENTITY_ID: lightb.entity_id,
-                CONF_MANUFACTURER: lightb.manufacturer,
-                CONF_MODEL: lightb.model,
-            },
-        ),
-    ]
+    mock_calls = mock_flow_init.mock_calls
+    assert len(mock_calls) == 2
+    assert mock_calls[0][2]["context"] == {"source": SOURCE_INTEGRATION_DISCOVERY}
+    assert mock_calls[0][2]["data"][CONF_ENTITY_ID] == "light.testa"
+    assert mock_calls[1][2]["context"] == {"source": SOURCE_INTEGRATION_DISCOVERY}
+    assert mock_calls[1][2]["data"][CONF_ENTITY_ID] == "light.testb"
 
     # Also check if power sensors are created.
     # Currently, we also create them directly, even without the user finishing the discovery flow
@@ -111,7 +93,7 @@ async def test_autodiscovery(hass: HomeAssistant, mock_flow_init):
 
 
 async def test_discovery_skipped_when_confirmed_by_user(
-    hass: HomeAssistant, mock_flow_init
+        hass: HomeAssistant, mock_flow_init
 ):
     light_entity = MockLight("test")
     light_entity.manufacturer = "lidl"
@@ -156,7 +138,7 @@ async def test_autodiscovery_disabled(hass: HomeAssistant):
 
 
 async def test_autodiscovery_skipped_for_lut_with_subprofiles(
-    hass: HomeAssistant, caplog: pytest.LogCaptureFixture
+        hass: HomeAssistant, caplog: pytest.LogCaptureFixture
 ):
     """
     Lights which can be autodiscovered and have sub profiles need to de skipped
@@ -196,7 +178,7 @@ async def test_manual_configured_light_overrides_autodiscovered(hass: HomeAssist
 
 
 async def test_config_entry_overrides_autodiscovered(
-    hass: HomeAssistant, caplog: pytest.LogCaptureFixture
+        hass: HomeAssistant, caplog: pytest.LogCaptureFixture
 ):
     caplog.set_level(logging.ERROR)