fix: tests and cleanup

custom_components/powercalc/config_flow.py

@@ -572,7 +572,7 @@ class Step(StrEnum):
 @dataclass(slots=True)
 class PowercalcFormStep:
     schema: vol.Schema | Callable[[], Coroutine[Any, Any, vol.Schema | None]]
-
+    step: Step
     validate_user_input: (
         Callable[
             [dict[str, Any]],
@@ -582,7 +582,6 @@ class PowercalcFormStep:
     ) = None
 
     next_step: Step | Callable[[dict[str, Any]], Coroutine[Any, Any, Step | None]] | None = None
-    step: Step | None = None
     continue_utility_meter_options_step: bool = False
     continue_advanced_step: bool = False
     form_kwarg: dict[str, Any] | None = None
@@ -602,6 +601,7 @@ def __init__(self) -> None:
         self.is_options_flow: bool = isinstance(self, OptionsFlow)
         self.strategy: CalculationStrategy | None = None
         self.name: str | None = None
+        self.handled_steps: list[Step] = []
         super().__init__()
 
     @abstractmethod
@@ -898,6 +898,7 @@ async def handle_form_step(
                 self.name = user_input[CONF_NAME]
             self.sensor_config.update(user_input)
 
+            self.handled_steps.append(form_step.step)
             next_step = form_step.next_step
             if callable(form_step.next_step):
                 next_step = await form_step.next_step(user_input)
@@ -1038,30 +1039,40 @@ async def async_step_post_library(
         if not self.selected_profile:
             return self.async_abort(reason="model_not_supported")  # pragma: no cover
 
-        if self.selected_profile.has_custom_fields and not self.sensor_config.get(CONF_VARIABLES):
+        if Step.LIBRARY_CUSTOM_FIELDS not in self.handled_steps and self.selected_profile.has_custom_fields:
             return await self.async_step_library_custom_fields()
 
-        if self.selected_profile.discovery_by == DiscoveryBy.DEVICE and not self.sensor_config.get(CONF_AVAILABILITY_ENTITY):
-            return await self.async_step_availability_entity()
+        if Step.AVAILABILITY_ENTITY not in self.handled_steps and self.selected_profile.discovery_by == DiscoveryBy.DEVICE:
+            result = await self.async_step_availability_entity()
+            if result:
+                return result
 
-        if await self.selected_profile.has_sub_profiles and not self.selected_profile.sub_profile_select:
+        if (
+            Step.SUB_PROFILE not in self.handled_steps
+            and await self.selected_profile.has_sub_profiles
+            and not self.selected_profile.sub_profile_select
+        ):
             return await self.async_step_sub_profile()
 
-        if self.selected_profile.device_type == DeviceType.SMART_SWITCH and self.selected_profile.calculation_strategy == CalculationStrategy.FIXED:
+        if (
+            Step.SMART_SWITCH not in self.handled_steps
+            and self.selected_profile.device_type == DeviceType.SMART_SWITCH
+            and self.selected_profile.calculation_strategy == CalculationStrategy.FIXED
+        ):
             return await self.async_step_smart_switch()
 
-        if self.selected_profile.needs_fixed_config:  # pragma: no cover
+        if Step.FIXED not in self.handled_steps and self.selected_profile.needs_fixed_config:  # pragma: no cover
             return await self.async_step_fixed()
 
-        if self.selected_profile.needs_linear_config:
+        if Step.LINEAR not in self.handled_steps and self.selected_profile.needs_linear_config:
             return await self.async_step_linear()
 
-        if self.selected_profile.calculation_strategy == CalculationStrategy.MULTI_SWITCH:
+        if Step.MULTI_SWITCH not in self.handled_steps and self.selected_profile.calculation_strategy == CalculationStrategy.MULTI_SWITCH:
             return await self.async_step_multi_switch()
 
         return await self.async_step_assign_groups()
 
-    async def async_step_availability_entity(self, user_input: dict[str, Any] | None = None) -> FlowResult:
+    async def async_step_availability_entity(self, user_input: dict[str, Any] | None = None) -> FlowResult | None:
         """Handle the flow for availability entity."""
         domains = DEVICE_TYPE_DOMAIN[self.selected_profile.device_type]  # type: ignore
         entity_selector = self.create_device_entity_selector(
@@ -1070,7 +1081,9 @@ async def async_step_availability_entity(self, user_input: dict[str, Any] | None
         try:
             first_entity = entity_selector.config["include_entities"][0]
         except IndexError:
-            first_entity = None
+            # Skip step if no entities are available
+            self.handled_steps.append(Step.AVAILABILITY_ENTITY)
+            return None
         return await self.handle_form_step(
             PowercalcFormStep(
                 step=Step.AVAILABILITY_ENTITY,

tests/config_flow/common.py

@@ -91,6 +91,18 @@ async def select_manufacturer_and_model(
     )
 
 
+async def confirm_auto_discovered_model(
+    hass: HomeAssistant,
+    prev_result: FlowResult,
+    confirmed: bool = True,
+) -> FlowResult:
+    assert prev_result["step_id"] == Step.LIBRARY
+    return await hass.config_entries.flow.async_configure(
+        prev_result["flow_id"],
+        {CONF_CONFIRM_AUTODISCOVERED_MODEL: confirmed},
+    )
+
+
 async def initialize_options_flow(
     hass: HomeAssistant,
     entry: config_entries.ConfigEntry,
@@ -155,11 +167,7 @@ async def initialize_discovery_flow(
     if not confirm_autodiscovered_model:
         return result
 
-    assert result["type"] == data_entry_flow.FlowResultType.FORM
-    return await hass.config_entries.flow.async_configure(
-        result["flow_id"],
-        {CONF_CONFIRM_AUTODISCOVERED_MODEL: True},
-    )
+    return await confirm_auto_discovered_model(hass, result)
 
 
 async def goto_virtual_power_strategy_step(

tests/config_flow/test_discovery.py

@@ -9,7 +9,7 @@
 from pytest_homeassistant_custom_component.common import mock_device_registry, mock_registry
 
 from custom_components.powercalc.common import SourceEntity, create_source_entity
-from custom_components.powercalc.config_flow import CONF_CONFIRM_AUTODISCOVERED_MODEL, Step
+from custom_components.powercalc.config_flow import Step
 from custom_components.powercalc.const import (
     CONF_AVAILABILITY_ENTITY,
     CONF_CREATE_ENERGY_SENSOR,
@@ -27,6 +27,7 @@
 from tests.config_flow.common import (
     DEFAULT_ENTITY_ID,
     DEFAULT_UNIQUE_ID,
+    confirm_auto_discovered_model,
     create_mock_entry,
     initialize_discovery_flow,
     initialize_options_flow,
@@ -49,11 +50,7 @@ async def test_discovery_flow(
     result = await initialize_discovery_flow(hass, source_entity)
 
     # Confirm selected manufacturer/model
-    assert result["type"] == data_entry_flow.FlowResultType.FORM
-    result = await hass.config_entries.flow.async_configure(
-        result["flow_id"],
-        {CONF_CONFIRM_AUTODISCOVERED_MODEL: True},
-    )
+    result = await confirm_auto_discovered_model(hass, result)
 
     assert result["type"] == data_entry_flow.FlowResultType.CREATE_ENTRY
     assert result["data"] == {
@@ -89,11 +86,7 @@ async def test_discovery_flow_with_subprofile_selection(
 
     result = await initialize_discovery_flow(hass, source_entity, power_profile)
 
-    assert result["type"] == data_entry_flow.FlowResultType.FORM
-    result = await hass.config_entries.flow.async_configure(
-        result["flow_id"],
-        {CONF_CONFIRM_AUTODISCOVERED_MODEL: True},
-    )
+    result = await confirm_auto_discovered_model(hass, result)
 
     assert result["type"] == data_entry_flow.FlowResultType.FORM
     assert result["step_id"] == Step.SUB_PROFILE
@@ -224,11 +217,7 @@ async def test_discovery_by_device(hass: HomeAssistant) -> None:
     ]
     result = await initialize_discovery_flow(hass, source_entity, power_profiles)
 
-    assert result["type"] == data_entry_flow.FlowResultType.FORM
-    result = await hass.config_entries.flow.async_configure(
-        result["flow_id"],
-        {CONF_CONFIRM_AUTODISCOVERED_MODEL: True},
-    )
+    result = await confirm_auto_discovered_model(hass, result)
 
     assert result["type"] == data_entry_flow.FlowResultType.FORM
     assert result["step_id"] == Step.AVAILABILITY_ENTITY

tests/config_flow/test_smart_switch.py

@@ -7,7 +7,7 @@
 from homeassistant.core import HomeAssistant
 
 from custom_components.powercalc import CONF_POWER
-from custom_components.powercalc.config_flow import CONF_CONFIRM_AUTODISCOVERED_MODEL, Step
+from custom_components.powercalc.config_flow import Step
 from custom_components.powercalc.const import (
     CONF_CREATE_ENERGY_SENSOR,
     CONF_CREATE_UTILITY_METERS,
@@ -23,6 +23,7 @@
 from tests.common import get_test_config_dir
 from tests.config_flow.common import (
     DEFAULT_UNIQUE_ID,
+    confirm_auto_discovered_model,
     create_mock_entry,
     initialize_options_flow,
     select_menu_item,
@@ -67,10 +68,7 @@ async def test_smart_switch_flow(
     assert result["type"] == data_entry_flow.FlowResultType.FORM
     assert result["step_id"] == Step.LIBRARY
 
-    result = await hass.config_entries.flow.async_configure(
-        result["flow_id"],
-        {CONF_CONFIRM_AUTODISCOVERED_MODEL: True},
-    )
+    result = await confirm_auto_discovered_model(hass, result)
 
     assert result["type"] == data_entry_flow.FlowResultType.FORM
     assert result["step_id"] == Step.SMART_SWITCH

tests/config_flow/test_virtual_power_library.py

@@ -5,14 +5,17 @@
 from homeassistant import data_entry_flow
 from homeassistant.const import CONF_ENTITY_ID, CONF_NAME, STATE_ON
 from homeassistant.core import HomeAssistant
+from homeassistant.helpers.device_registry import DeviceEntry
 from homeassistant.helpers.selector import SelectSelector
+from pytest_homeassistant_custom_component.common import mock_device_registry
 
 from custom_components.powercalc import (
     CONF_CREATE_ENERGY_SENSOR,
     CONF_CREATE_UTILITY_METERS,
     CONF_ENERGY_INTEGRATION_METHOD,
     DEFAULT_ENERGY_INTEGRATION_METHOD,
 )
+from custom_components.powercalc.common import create_source_entity
 from custom_components.powercalc.config_flow import (
     CONF_CONFIRM_AUTODISCOVERED_MODEL,
     Step,
@@ -28,12 +31,16 @@
     CalculationStrategy,
     SensorType,
 )
+from custom_components.powercalc.power_profile.factory import get_power_profile
+from custom_components.powercalc.power_profile.library import ModelInfo
 from custom_components.test.light import MockLight
 from tests.common import create_mock_light_entity, get_test_config_dir
 from tests.config_flow.common import (
     DEFAULT_UNIQUE_ID,
+    confirm_auto_discovered_model,
     create_mock_entry,
     goto_virtual_power_strategy_step,
+    initialize_discovery_flow,
     initialize_options_flow,
     process_config_flow,
     select_manufacturer_and_model,
@@ -351,3 +358,25 @@ async def test_sub_profiles_select_options(hass: HomeAssistant) -> None:
         {},
     )
     assert result["type"] == data_entry_flow.FlowResultType.CREATE_ENTRY
+
+
+async def test_availability_entity_step_skipped(hass: HomeAssistant) -> None:
+    hass.config.config_dir = get_test_config_dir()
+    mock_device_registry(
+        hass,
+        {
+            "test-device": DeviceEntry(
+                manufacturer="test",
+                name="Test Device",
+                model="discovery_type_device",
+            ),
+        },
+    )
+
+    source_entity = await create_source_entity(DUMMY_ENTITY_ID, hass)
+    power_profiles = [
+        await get_power_profile(hass, {}, ModelInfo("test", "discovery_type_device")),
+    ]
+    result = await initialize_discovery_flow(hass, source_entity, power_profiles)
+    result = await confirm_auto_discovered_model(hass, result)
+    assert result["type"] == data_entry_flow.FlowResultType.CREATE_ENTRY

tests/config_flow/test_virtual_power_lut.py

@@ -21,6 +21,7 @@
 from tests.config_flow.common import (
     DEFAULT_UNIQUE_ID,
     assert_default_virtual_power_entry_data,
+    confirm_auto_discovered_model,
     create_mock_entry,
     goto_virtual_power_strategy_step,
     initialize_options_flow,
@@ -160,10 +161,7 @@ async def test_lut_autodiscover_flow_not_confirmed(
     assert result["type"] == data_entry_flow.FlowResultType.FORM
     assert result["step_id"] == Step.LIBRARY
 
-    result = await hass.config_entries.flow.async_configure(
-        result["flow_id"],
-        {CONF_CONFIRM_AUTODISCOVERED_MODEL: False},
-    )
+    result = await confirm_auto_discovered_model(hass, result, confirmed=False)
 
     assert result["type"] == data_entry_flow.FlowResultType.FORM
     assert result["step_id"] == Step.MANUFACTURER

tests/power_profile/device_types/test_smart_dimmer.py

@@ -4,7 +4,7 @@
 from homeassistant.data_entry_flow import FlowResultType
 
 from custom_components.powercalc import CONF_IGNORE_UNAVAILABLE_STATE, async_setup_entry
-from custom_components.powercalc.config_flow import CONF_CONFIRM_AUTODISCOVERED_MODEL, Step
+from custom_components.powercalc.config_flow import Step
 from custom_components.powercalc.const import (
     CONF_CUSTOM_MODEL_DIRECTORY,
     CONF_LINEAR,
@@ -13,7 +13,7 @@
     DOMAIN,
 )
 from tests.common import get_test_config_dir, get_test_profile_dir, run_powercalc_setup
-from tests.config_flow.common import initialize_options_flow
+from tests.config_flow.common import confirm_auto_discovered_model, initialize_options_flow
 from tests.conftest import MockEntityWithModel
 
 
@@ -115,10 +115,7 @@ async def test_smart_dimmer_power_input_gui_config_flow(
     flow = flows[0]
 
     assert flow["step_id"] == Step.LIBRARY
-    result = await hass.config_entries.flow.async_configure(
-        flow["flow_id"],
-        {CONF_CONFIRM_AUTODISCOVERED_MODEL: True},
-    )
+    result = await confirm_auto_discovered_model(hass, flow)
 
     # After confirming the manufacturer/model we must be directed to the linear config step
     assert result["step_id"] == Step.LINEAR

tests/power_profile/device_types/test_smart_switch.py

@@ -6,7 +6,7 @@
 from homeassistant.data_entry_flow import FlowResult, FlowResultType
 
 from custom_components.powercalc import async_setup_entry
-from custom_components.powercalc.config_flow import CONF_CONFIRM_AUTODISCOVERED_MODEL, Step
+from custom_components.powercalc.config_flow import Step
 from custom_components.powercalc.const import (
     CONF_CUSTOM_MODEL_DIRECTORY,
     CONF_FIXED,
@@ -16,7 +16,7 @@
     DOMAIN,
 )
 from tests.common import get_test_config_dir, get_test_profile_dir, run_powercalc_setup
-from tests.config_flow.common import initialize_options_flow
+from tests.config_flow.common import confirm_auto_discovered_model, initialize_options_flow
 from tests.conftest import MockEntityWithModel
 
 
@@ -243,7 +243,4 @@ async def start_discovery_flow(
     flow = flows[0]
 
     assert flow["step_id"] == Step.LIBRARY
-    return await hass.config_entries.flow.async_configure(
-        flow["flow_id"],
-        {CONF_CONFIRM_AUTODISCOVERED_MODEL: True},
-    )
+    return await confirm_auto_discovered_model(hass, flow)

tests/power_profile/user_scenarios/test_lightify_plug.py

@@ -4,11 +4,11 @@
 from homeassistant.const import CONF_ENTITY_ID
 from homeassistant.core import HomeAssistant
 
-from custom_components.powercalc.config_flow import CONF_CONFIRM_AUTODISCOVERED_MODEL, Step
+from custom_components.powercalc.config_flow import Step
 from custom_components.powercalc.const import (
     CONF_MANUFACTURER,
 )
-from tests.config_flow.common import select_menu_item
+from tests.config_flow.common import confirm_auto_discovered_model, select_menu_item
 from tests.conftest import MockEntityWithModel
 
 
@@ -31,10 +31,7 @@ async def test_lightify_plug_selectable(
         result["flow_id"],
         {CONF_ENTITY_ID: "light.test"},
     )
-    result = await hass.config_entries.flow.async_configure(
-        result["flow_id"],
-        {CONF_CONFIRM_AUTODISCOVERED_MODEL: False},
-    )
+    result = await confirm_auto_discovered_model(hass, result, confirmed=False)
     result = await hass.config_entries.flow.async_configure(
         result["flow_id"],
         {