Add tests

fix some bugs
marcolivierarsenault · Jul 9, 2023 · 57c0ed2 · 57c0ed2
1 parent 55be85d
commit 57c0ed2
Show file tree

Hide file tree

Showing 2 changed files with 53 additions and 2 deletions.
diff --git a/custom_components/moonraker/sensor.py b/custom_components/moonraker/sensor.py
@@ -524,7 +524,11 @@ def calculate_current_layer(data):
     ):
         return 0
 
-    if (data["status"]["print_stats"]["info"]["current_layer"] != None):
+    if (
+        "info" in data["status"]["print_stats"]
+        and "current_layer" in data["status"]["print_stats"]["info"]
+        and data["status"]["print_stats"]["info"]["current_layer"] is not None
+    ):
         return data["status"]["print_stats"]["info"]["current_layer"]
 
     # layer = (current_z - first_layer_height) / layer_height + 1

diff --git a/tests/test_sensor.py b/tests/test_sensor.py
@@ -11,7 +11,10 @@
 
 from custom_components.moonraker import async_setup_entry
 from custom_components.moonraker.const import DOMAIN, PRINTSTATES
-from custom_components.moonraker.sensor import calculate_pct_job
+from custom_components.moonraker.sensor import (
+    calculate_current_layer,
+    calculate_pct_job,
+)
 
 from .const import MOCK_CONFIG
 
@@ -285,3 +288,47 @@ async def test_rounding_fan(hass, get_data):
 
     state = hass.states.get("sensor.mainsail_fan_speed")
     assert state.state == "33.33"
+
+
+async def test_current_layer_in_info():
+    data = {
+        "status": {
+            "print_stats": {
+                "state": PRINTSTATES.PRINTING.value,
+                "filename": "TheUniverse.gcode",
+                "info": {"current_layer": 42},
+            },
+        },
+    }
+    assert calculate_current_layer(data) == 42
+
+
+async def test_current_layer_calculated():
+    data = {
+        "status": {
+            "print_stats": {
+                "state": PRINTSTATES.PRINTING.value,
+                "filename": "TheUniverse.gcode",
+            },
+            "toolhead": {"position": [0, 0, 8.4]},
+        },
+        "first_layer_height": 0.2,
+        "layer_height": 0.2,
+    }
+    assert calculate_current_layer(data) == 42
+
+
+async def test_current_layer_calculated_partial_info():
+    data = {
+        "status": {
+            "print_stats": {
+                "state": PRINTSTATES.PRINTING.value,
+                "filename": "TheUniverse.gcode",
+                "info": {},
+            },
+            "toolhead": {"position": [0, 0, 8.4]},
+        },
+        "first_layer_height": 0.2,
+        "layer_height": 0.2,
+    }
+    assert calculate_current_layer(data) == 42