From 075152cb1709b01426f2c538199dcc61a6e6a0ab Mon Sep 17 00:00:00 2001
From: Bram Gerritsen <bgerritsen@emico.nl>
Date: Sat, 26 Oct 2024 08:48:27 +0200
Subject: [PATCH] fix: multi switch setup using GUI (#2623)

---
 custom_components/powercalc/const.py          |  5 +++
 custom_components/powercalc/sensors/power.py  | 10 ++---
 .../test_virtual_power_multi_switch.py        | 37 ++++++++++++++++++-
 3 files changed, 44 insertions(+), 8 deletions(-)

diff --git a/custom_components/powercalc/const.py b/custom_components/powercalc/const.py
index db3470778..701071cb1 100644
--- a/custom_components/powercalc/const.py
+++ b/custom_components/powercalc/const.py
@@ -1,5 +1,7 @@
 """The Powercalc constants."""
 
+from __future__ import annotations
+
 from datetime import timedelta
 from enum import StrEnum
 from typing import Literal
@@ -207,6 +209,9 @@ class CalculationStrategy(StrEnum):
     WLED = "wled"
 
 
+CALCULATION_STRATEGY_CONF_KEYS: list[str] = [strategy.value for strategy in CalculationStrategy]
+
+
 class SensorType(StrEnum):
     """Possible modes for a number selector."""
 
diff --git a/custom_components/powercalc/sensors/power.py b/custom_components/powercalc/sensors/power.py
index 03122fb3e..775a63aa9 100644
--- a/custom_components/powercalc/sensors/power.py
+++ b/custom_components/powercalc/sensors/power.py
@@ -50,19 +50,16 @@
     ATTR_INTEGRATION,
     ATTR_SOURCE_DOMAIN,
     ATTR_SOURCE_ENTITY,
+    CALCULATION_STRATEGY_CONF_KEYS,
     CONF_CALCULATION_ENABLED_CONDITION,
-    CONF_COMPOSITE,
     CONF_DELAY,
     CONF_DISABLE_EXTENDED_ATTRIBUTES,
     CONF_DISABLE_STANDBY_POWER,
-    CONF_FIXED,
     CONF_FORCE_UPDATE_FREQUENCY,
     CONF_IGNORE_UNAVAILABLE_STATE,
-    CONF_LINEAR,
     CONF_MODEL,
     CONF_MULTIPLY_FACTOR,
     CONF_MULTIPLY_FACTOR_STANDBY,
-    CONF_PLAYBOOK,
     CONF_POWER,
     CONF_POWER_SENSOR_CATEGORY,
     CONF_POWER_SENSOR_ID,
@@ -70,7 +67,6 @@
     CONF_SLEEP_POWER,
     CONF_STANDBY_POWER,
     CONF_UNAVAILABLE_POWER,
-    CONF_WLED,
     DATA_CALCULATOR_FACTORY,
     DATA_DISCOVERY_MANAGER,
     DATA_STANDBY_POWER_SENSORS,
@@ -299,11 +295,11 @@ def is_manually_configured(sensor_config: ConfigType) -> bool:
     """
     if CONF_MODEL in sensor_config:
         return False
-    return any(key in sensor_config for key in [CONF_LINEAR, CONF_FIXED, CONF_PLAYBOOK, CONF_COMPOSITE])
+    return any(key in sensor_config for key in CALCULATION_STRATEGY_CONF_KEYS)
 
 
 def is_fully_configured(config: ConfigType) -> bool:
-    return any(key in config for key in [CONF_LINEAR, CONF_WLED, CONF_FIXED, CONF_PLAYBOOK])
+    return any(key in config for key in CALCULATION_STRATEGY_CONF_KEYS)
 
 
 class PowerSensor(BaseEntity):
diff --git a/tests/config_flow/test_virtual_power_multi_switch.py b/tests/config_flow/test_virtual_power_multi_switch.py
index b1dac8786..1dde6a2a9 100644
--- a/tests/config_flow/test_virtual_power_multi_switch.py
+++ b/tests/config_flow/test_virtual_power_multi_switch.py
@@ -22,7 +22,7 @@
     CalculationStrategy,
     SensorType,
 )
-from tests.common import get_test_config_dir
+from tests.common import get_test_config_dir, run_powercalc_setup
 from tests.config_flow.common import (
     DEFAULT_UNIQUE_ID,
     create_mock_entry,
@@ -170,3 +170,38 @@ async def test_options_flow(hass: HomeAssistant) -> None:
 
     assert result["type"] == data_entry_flow.FlowResultType.CREATE_ENTRY
     assert entry.data[CONF_MULTI_SWITCH] == {CONF_POWER: 5, CONF_POWER_OFF: 20, CONF_ENTITIES: ["switch.a", "switch.c"]}
+
+
+async def test_regression_2612(hass: HomeAssistant, mock_entity_with_model_information: MockEntityWithModel) -> None:
+    """
+    See #2612
+    When the source entity had manufacturer and model information the multi switch setup would fail
+    And raise error "Model not found in library" in the logs
+    """
+
+    mock_entity_with_model_information(
+        "switch.test",
+        "_TZ3000_u3oupgdy",
+        "TS0004",
+        unique_id=DEFAULT_UNIQUE_ID,
+    )
+
+    create_mock_entry(
+        hass,
+        {
+            CONF_ENTITY_ID: "switch.test",
+            CONF_SENSOR_TYPE: SensorType.VIRTUAL_POWER,
+            CONF_MODE: CalculationStrategy.MULTI_SWITCH,
+            CONF_MULTI_SWITCH: {
+                CONF_POWER: 10,
+                CONF_POWER_OFF: 40,
+                CONF_ENTITIES: ["switch.a", "switch.b"],
+            },
+            CONF_NAME: "Foo bar",
+        },
+    )
+
+    await run_powercalc_setup(hass, {})
+
+    assert hass.states.get("sensor.foo_bar_power")
+    assert hass.states.get("sensor.foo_bar_energy")