From 2847524b2168fb21afa556d9276829ac8d77f0a0 Mon Sep 17 00:00:00 2001
From: Marc-Olivier Arsenault <marcolivier.arsenault@gmail.com>
Date: Fri, 16 Aug 2024 11:31:57 +0000
Subject: [PATCH] fix sensor, account for missing data

---
 custom_components/moonraker/sensor.py | 38 +++++++++++++++++----------
 tests/test_sensor.py                  |  2 +-
 2 files changed, 25 insertions(+), 15 deletions(-)

diff --git a/custom_components/moonraker/sensor.py b/custom_components/moonraker/sensor.py
index 491f123..ea2b8c5 100755
--- a/custom_components/moonraker/sensor.py
+++ b/custom_components/moonraker/sensor.py
@@ -415,16 +415,21 @@ async def async_setup_optional_sensors(coordinator, entry, async_add_entities):
                 status_key=obj,
                 name=f"{name} Load",
                 value_fn=lambda sensor: (
-                    sensor.coordinator.data["status"][sensor.status_key]["last_stats"][
-                        "mcu_task_avg"
-                    ]
-                    + 3
-                    * sensor.coordinator.data["status"][sensor.status_key][
-                        "last_stats"
-                    ]["mcu_task_stddev"]
+                    (
+                        sensor.coordinator.data["status"][sensor.status_key][
+                            "last_stats"
+                        ]["mcu_task_avg"]
+                        + 3
+                        * sensor.coordinator.data["status"][sensor.status_key][
+                            "last_stats"
+                        ]["mcu_task_stddev"]
+                    )
+                    / 0.0025
+                    * 100
                 )
-                / 0.0025
-                * 100,
+                if sensor.coordinator.data["status"][sensor.status_key]["last_stats"]
+                is not None
+                else 0,
                 subscriptions=[(obj, "last_stats")],
                 icon="mdi:cpu-64-bit",
                 state_class=SensorStateClass.MEASUREMENT,
@@ -435,11 +440,16 @@ async def async_setup_optional_sensors(coordinator, entry, async_add_entities):
                 key=f"{key}_awake",
                 status_key=obj,
                 name=f"{name} Awake",
-                value_fn=lambda sensor: sensor.coordinator.data["status"][
-                    sensor.status_key
-                ]["last_stats"]["mcu_awake"]
-                / 5
-                * 100,
+                value_fn=lambda sensor: (
+                    sensor.coordinator.data["status"][sensor.status_key]["last_stats"][
+                        "mcu_awake"
+                    ]
+                    / 5
+                    * 100
+                )
+                if sensor.coordinator.data["status"][sensor.status_key]["last_stats"]
+                is not None
+                else 0,
                 icon="mdi:cpu-64-bit",
                 subscriptions=[(obj, "last_stats")],
                 state_class=SensorStateClass.MEASUREMENT,
diff --git a/tests/test_sensor.py b/tests/test_sensor.py
index f421f30..5f3022e 100755
--- a/tests/test_sensor.py
+++ b/tests/test_sensor.py
@@ -374,7 +374,7 @@ async def test_multi_mcu_sensor_data(hass, get_data, get_printer_objects_list):
 
     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.config_entries.async_setup(config_entry.entry_id)
     await hass.async_block_till_done()
 
     registry = get_entity_registry(hass)