From a25653e167dccc47236fc3c6dd7cd924010f9eb4 Mon Sep 17 00:00:00 2001
From: Josef Zweck <24647999+zweckj@users.noreply.github.com>
Date: Wed, 17 Jan 2024 11:07:43 +0100
Subject: [PATCH] Change the way an entity is supported in La Marzocco
 (#108216)

* refactor supported

* refactor supported
---
 homeassistant/components/lamarzocco/entity.py |  8 +-------
 homeassistant/components/lamarzocco/sensor.py | 17 ++++++-----------
 homeassistant/components/lamarzocco/switch.py |  2 +-
 3 files changed, 8 insertions(+), 19 deletions(-)

diff --git a/homeassistant/components/lamarzocco/entity.py b/homeassistant/components/lamarzocco/entity.py
index 4b8ccc86688..6918741f1d3 100644
--- a/homeassistant/components/lamarzocco/entity.py
+++ b/homeassistant/components/lamarzocco/entity.py
@@ -4,7 +4,6 @@ from collections.abc import Callable
 from dataclasses import dataclass
 
 from lmcloud import LMCloud as LaMarzoccoClient
-from lmcloud.const import LaMarzoccoModel
 
 from homeassistant.helpers.device_registry import DeviceInfo
 from homeassistant.helpers.entity import EntityDescription
@@ -19,12 +18,7 @@ class LaMarzoccoEntityDescription(EntityDescription):
     """Description for all LM entities."""
 
     available_fn: Callable[[LaMarzoccoClient], bool] = lambda _: True
-    supported_models: tuple[LaMarzoccoModel, ...] = (
-        LaMarzoccoModel.GS3_AV,
-        LaMarzoccoModel.GS3_MP,
-        LaMarzoccoModel.LINEA_MICRA,
-        LaMarzoccoModel.LINEA_MINI,
-    )
+    supported_fn: Callable[[LaMarzoccoUpdateCoordinator], bool] = lambda _: True
 
 
 class LaMarzoccoEntity(CoordinatorEntity[LaMarzoccoUpdateCoordinator]):
diff --git a/homeassistant/components/lamarzocco/sensor.py b/homeassistant/components/lamarzocco/sensor.py
index bb811eaa890..63292b95ae3 100644
--- a/homeassistant/components/lamarzocco/sensor.py
+++ b/homeassistant/components/lamarzocco/sensor.py
@@ -59,6 +59,7 @@ ENTITIES: tuple[LaMarzoccoSensorEntityDescription, ...] = (
         value_fn=lambda lm: lm.current_status.get("brew_active_duration", 0),
         available_fn=lambda lm: lm.websocket_connected,
         entity_category=EntityCategory.DIAGNOSTIC,
+        supported_fn=lambda coordinator: coordinator.local_connection_configured,
     ),
     LaMarzoccoSensorEntityDescription(
         key="current_temp_coffee",
@@ -89,17 +90,11 @@ async def async_setup_entry(
     """Set up sensor entities."""
     coordinator = hass.data[DOMAIN][config_entry.entry_id]
 
-    entities: list[LaMarzoccoSensorEntity] = []
-    for description in ENTITIES:
-        if coordinator.lm.model_name in description.supported_models:
-            if (
-                description.key == "shot_timer"
-                and not coordinator.local_connection_configured
-            ):
-                continue
-            entities.append(LaMarzoccoSensorEntity(coordinator, description))
-
-    async_add_entities(entities)
+    async_add_entities(
+        LaMarzoccoSensorEntity(coordinator, description)
+        for description in ENTITIES
+        if description.supported_fn(coordinator)
+    )
 
 
 class LaMarzoccoSensorEntity(LaMarzoccoEntity, SensorEntity):
diff --git a/homeassistant/components/lamarzocco/switch.py b/homeassistant/components/lamarzocco/switch.py
index 4c39bd2c5f0..fe9c6daa9cf 100644
--- a/homeassistant/components/lamarzocco/switch.py
+++ b/homeassistant/components/lamarzocco/switch.py
@@ -67,7 +67,7 @@ async def async_setup_entry(
     async_add_entities(
         LaMarzoccoSwitchEntity(coordinator, description)
         for description in ENTITIES
-        if coordinator.lm.model_name in description.supported_models
+        if description.supported_fn(coordinator)
     )