From 09f250d7fe2ce0288bf8ae2306771495afde70f3 Mon Sep 17 00:00:00 2001
From: Eric Tremblay <etremblay@distech-controls.com>
Date: Sun, 21 Apr 2024 19:54:34 -0400
Subject: [PATCH] Make the heater_bed printer object optional

The patch adds the bed sensors (temperature, target, power) only when
the heater_bed object is available.

This fix cases where the integration was not starting without the
heater_bed object.
---
 custom_components/moonraker/sensor.py | 52 +++++++--------------------
 tests/test_sensor.py                  | 12 -------
 2 files changed, 13 insertions(+), 51 deletions(-)

diff --git a/custom_components/moonraker/sensor.py b/custom_components/moonraker/sensor.py
index 4c8bb81..ef47609 100755
--- a/custom_components/moonraker/sensor.py
+++ b/custom_components/moonraker/sensor.py
@@ -84,28 +84,6 @@ class MoonrakerSensorDescription(SensorEntityDescription):
         else "",
         subscriptions=[("display_status", "message")],
     ),
-    MoonrakerSensorDescription(
-        key="bed_target",
-        name="Bed Target",
-        value_fn=lambda sensor: float(
-            sensor.coordinator.data["status"]["heater_bed"]["target"] or 0.0
-        ),
-        subscriptions=[("heater_bed", "target")],
-        icon="mdi:radiator",
-        unit=UnitOfTemperature.CELSIUS,
-        state_class=SensorStateClass.MEASUREMENT,
-    ),
-    MoonrakerSensorDescription(
-        key="bed_temp",
-        name="Bed Temperature",
-        value_fn=lambda sensor: float(
-            sensor.coordinator.data["status"]["heater_bed"]["temperature"] or 0.0
-        ),
-        subscriptions=[("heater_bed", "temperature")],
-        icon="mdi:radiator",
-        unit=UnitOfTemperature.CELSIUS,
-        state_class=SensorStateClass.MEASUREMENT,
-    ),
     MoonrakerSensorDescription(
         key="filename",
         name="Filename",
@@ -230,17 +208,6 @@ class MoonrakerSensorDescription(SensorEntityDescription):
         icon="mdi:percent",
         unit=PERCENTAGE,
     ),
-    MoonrakerSensorDescription(
-        key="bed_power",
-        name="Bed Power",
-        value_fn=lambda sensor: int(
-            sensor.coordinator.data["status"]["heater_bed"]["power"] * 100
-        ),
-        subscriptions=[("heater_bed", "power")],
-        icon="mdi:flash",
-        unit=PERCENTAGE,
-        state_class=SensorStateClass.MEASUREMENT,
-    ),
     MoonrakerSensorDescription(
         key="total_layer",
         name="Total Layer",
@@ -492,17 +459,24 @@ async def async_setup_optional_sensors(coordinator, entry, async_add_entities):
                 state_class=SensorStateClass.MEASUREMENT,
             )
             sensors.append(desc)
-        elif obj.startswith("extruder"):
+        elif obj.startswith("extruder") or obj.startswith("heater_bed"):
+            if obj.startswith("extruder"):
+                icon = "mdi:printer-3d-nozzle-heat"
+                base_name = obj
+            else:
+                icon = "mdi:radiator"
+                base_name = "Bed"
+
             desc = MoonrakerSensorDescription(
                 key=f"{obj}_temp",
                 status_key=obj,
-                name=f"{obj} Temperature".title(),
+                name=f"{base_name} Temperature".title(),
                 value_fn=lambda sensor: float(
                     sensor.coordinator.data["status"][sensor.status_key]["temperature"]
                     or 0.0
                 ),
                 subscriptions=[(obj, "temperature")],
-                icon="mdi:printer-3d-nozzle-heat",
+                icon=icon,
                 unit=UnitOfTemperature.CELSIUS,
                 state_class=SensorStateClass.MEASUREMENT,
             )
@@ -511,13 +485,13 @@ async def async_setup_optional_sensors(coordinator, entry, async_add_entities):
             desc = MoonrakerSensorDescription(
                 key=f"{obj}_target",
                 status_key=obj,
-                name=f"{obj} Target".title(),
+                name=f"{base_name} Target".title(),
                 value_fn=lambda sensor: float(
                     sensor.coordinator.data["status"][sensor.status_key]["target"]
                     or 0.0
                 ),
                 subscriptions=[(obj, "target")],
-                icon="mdi:printer-3d-nozzle-heat",
+                icon=icon,
                 unit=UnitOfTemperature.CELSIUS,
                 state_class=SensorStateClass.MEASUREMENT,
             )
@@ -526,7 +500,7 @@ async def async_setup_optional_sensors(coordinator, entry, async_add_entities):
             desc = MoonrakerSensorDescription(
                 key=f"{obj}_power",
                 status_key=obj,
-                name=f"{obj} Power".title(),
+                name=f"{base_name} Power".title(),
                 value_fn=lambda sensor: int(
                     sensor.coordinator.data["status"][sensor.status_key]["power"] * 100
                 ),
diff --git a/tests/test_sensor.py b/tests/test_sensor.py
index 7d7b1c8..3609af4 100755
--- a/tests/test_sensor.py
+++ b/tests/test_sensor.py
@@ -216,18 +216,6 @@ async def test_opt_sensor_missing(hass, get_data, get_printer_objects_list):
     assert state is None
 
 
-async def test_missing_heater_bed(hass, get_data):
-    """Test."""
-    get_data["status"]["heater_bed"]["target"] = None
-
-    config_entry = MockConfigEntry(domain=DOMAIN, data=MOCK_CONFIG, entry_id="test")
-    config_entry.add_to_hass(hass)
-    assert await async_setup_entry(hass, config_entry)
-    await hass.async_block_till_done()
-
-    assert hass.states.get("sensor.mainsail_bed_target").state == "0.0"
-
-
 async def test_eta(hass):
     """Test."""
     config_entry = MockConfigEntry(domain=DOMAIN, data=MOCK_CONFIG, entry_id="test")