feat: bambu-watch extensions — temp thresholds (#9) + heartbeat (#14)

src/bambu_ai/watcher.py

@@ -1,13 +1,18 @@
-"""Long-running watcher — pushes ntfy.sh notifications on print state transitions and HMS alerts.
+"""Long-running watcher — pushes ntfy.sh notifications on print state transitions,
+HMS alerts, optional temperature-threshold reaches, and optional heartbeat (printer
+unreachable) alerts.
 
 Entry point: ``bambu-watch`` (defined in pyproject.toml ``[project.scripts]``).
 """
 
 from __future__ import annotations
 
+import argparse
 import json
 import ssl
 import time
+from collections.abc import Callable
+from dataclasses import dataclass
 
 import paho.mqtt.client as mqtt
 
@@ -25,6 +30,9 @@
 ]
 
 
+# --------------------------------------------------------------------- helpers
+
+
 def maybe_notify_state(cfg: dict, prev: str | None, curr: str, snap: dict) -> None:
     file = snap.get("subtask_name") or snap.get("gcode_file") or "?"
     for from_s, to_s, title, body_t, prio, tags in RULES:
@@ -39,14 +47,153 @@ def hms_key(hms: list) -> str:
     return ", ".join(f"{h.get('attr', 0):08X}:{h.get('code', 0):08X}" for h in hms)
 
 
+# ----------------------------------------------------- temperature thresholds (#9)
+
+
+@dataclass
+class ThresholdState:
+    """Tracks whether configured bed/nozzle thresholds have been crossed during
+    the current print. ``check_and_notify`` fires once per crossing, then
+    re-arms on the next IDLE/PREPARE -> RUNNING transition (see
+    ``reset_for_new_print``)."""
+
+    bed_threshold: float | None = None
+    nozzle_threshold: float | None = None
+    bed_reached: bool = False
+    nozzle_reached: bool = False
+
+    def reset_for_new_print(self) -> None:
+        self.bed_reached = False
+        self.nozzle_reached = False
+
+    def check_and_notify(self, snap: dict, notifier: Callable[[str, str, str, str, str], None]) -> None:
+        """Inspect the snapshot for current temps; fire one notification per
+        crossing. ``notifier`` matches the :func:`notify` signature
+        ``(topic, title, body, priority, tags)`` so tests can inject a mock."""
+        if self.bed_threshold is not None and not self.bed_reached:
+            t = snap.get("bed_temper")
+            if isinstance(t, (int, float)) and t >= self.bed_threshold:
+                self.bed_reached = True
+                notifier(
+                    "",  # topic injected by caller
+                    "Bed reached target",
+                    f"Bed is at {t:.0f}°C (≥ {self.bed_threshold:.0f}°C target).",
+                    "default",
+                    "thermometer",
+                )
+        if self.nozzle_threshold is not None and not self.nozzle_reached:
+            t = snap.get("nozzle_temper")
+            if isinstance(t, (int, float)) and t >= self.nozzle_threshold:
+                self.nozzle_reached = True
+                notifier(
+                    "",
+                    "Nozzle reached target",
+                    f"Nozzle is at {t:.0f}°C (≥ {self.nozzle_threshold:.0f}°C target).",
+                    "default",
+                    "thermometer",
+                )
+
+
+# ----------------------------------------------------------------- heartbeat (#14)
+
+
+@dataclass
+class HeartbeatState:
+    """When enabled, fires one ntfy alert on each connection loss and one
+    'back online' alert on each successful reconnect."""
+
+    enabled: bool = False
+    last_disconnect_rc: int | None = None
+    disconnect_count: int = 0
+    in_alarm: bool = False
+
+    def on_disconnect(
+        self,
+        rc: int,
+        notifier: Callable[[str, str, str, str, str], None],
+    ) -> None:
+        self.disconnect_count += 1
+        self.last_disconnect_rc = rc
+        if not self.enabled or self.in_alarm:
+            return
+        # rc == 0 means clean local disconnect (e.g. Ctrl-C). Don't fire on that.
+        if rc == 0:
+            return
+        self.in_alarm = True
+        notifier(
+            "",
+            "Printer unreachable",
+            f"Lost MQTT connection (code {rc}). The watcher will keep retrying.",
+            "high",
+            "warning",
+        )
+
+    def on_connect(self, notifier: Callable[[str, str, str, str, str], None]) -> None:
+        if not self.enabled:
+            return
+        if self.in_alarm:
+            self.in_alarm = False
+            notifier(
+                "",
+                "Printer back online",
+                f"Reconnected after {self.disconnect_count} disconnect(s).",
+                "default",
+                "white_check_mark",
+            )
+
+
+# ----------------------------------------------------------------------- argparse
+
+
+def _parse_args(argv: list[str] | None = None) -> argparse.Namespace:
+    p = argparse.ArgumentParser(
+        prog="bambu-watch",
+        description="Long-running notifier for state transitions, HMS errors, temperature reaches, and printer connectivity.",
+    )
+    p.add_argument(
+        "--notify-bed-temp",
+        type=float,
+        default=None,
+        metavar="C",
+        help="Notify once when bed temperature crosses this threshold (°C). Re-arms on next print.",
+    )
+    p.add_argument(
+        "--notify-nozzle-temp",
+        type=float,
+        default=None,
+        metavar="C",
+        help="Notify once when nozzle temperature crosses this threshold (°C). Re-arms on next print.",
+    )
+    p.add_argument(
+        "--heartbeat",
+        action="store_true",
+        help="Fire ntfy alerts on MQTT disconnect/reconnect (catches overnight WiFi-dropped class of bugs).",
+    )
+    return p.parse_args(argv)
+
+
+# ----------------------------------------------------------------------- main
+
+
 def main() -> None:
+    args = _parse_args()
     cfg = load_config(extra_required=("NTFY_TOPIC",))
     topic_report = f"device/{cfg['SERIAL']}/report"
     topic_request = f"device/{cfg['SERIAL']}/request"
     snap: dict = {}
     seen_hms: set[str] = set()
     last_state: str | None = None
 
+    thresholds = ThresholdState(
+        bed_threshold=args.notify_bed_temp,
+        nozzle_threshold=args.notify_nozzle_temp,
+    )
+    heartbeat = HeartbeatState(enabled=args.heartbeat)
+
+    def _notify(_topic_unused: str, title: str, body: str, priority: str, tags: str) -> None:
+        # Helper threshold/heartbeat methods don't know the user's topic — bind it here.
+        notify(cfg["NTFY_TOPIC"], title, body, priority, tags)
+
     client = mqtt.Client(
         mqtt.CallbackAPIVersion.VERSION2,
         client_id=f"bambu-watcher-{int(time.time())}",
@@ -58,11 +205,16 @@ def main() -> None:
     def on_connect(c, _u, _f, rc, _p=None):
         if rc == 0:
             print(f"watching {cfg['IP']} — ntfy topic: {cfg['NTFY_TOPIC']}")
+            if thresholds.bed_threshold or thresholds.nozzle_threshold:
+                print(f"  threshold: bed={thresholds.bed_threshold}, nozzle={thresholds.nozzle_threshold}")
+            if heartbeat.enabled:
+                print("  heartbeat: enabled")
             c.subscribe(topic_report)
             c.publish(
                 topic_request,
                 json.dumps({"pushing": {"sequence_id": "0", "command": "pushall"}}),
             )
+            heartbeat.on_connect(_notify)
         else:
             print(f"connect failed: {rc}")
 
@@ -82,8 +234,13 @@ def on_message(_c, _u, msg):
             print(f"[state] {last_state} -> {curr}")
             if last_state is not None:
                 maybe_notify_state(cfg, last_state, curr, snap)
+            # Re-arm temperature thresholds on print start.
+            if curr == "RUNNING" and last_state in ("IDLE", "PREPARE", None):
+                thresholds.reset_for_new_print()
             last_state = curr
 
+        thresholds.check_and_notify(snap, _notify)
+
         hms = p.get("hms") or []
         if hms:
             key = hms_key(hms)
@@ -100,6 +257,7 @@ def on_message(_c, _u, msg):
 
     def on_disconnect(_c, _u, _f, rc, _p=None):
         print(f"disconnected: {rc} — will auto-reconnect")
+        heartbeat.on_disconnect(rc, _notify)
 
     client.on_connect = on_connect
     client.on_message = on_message

tests/test_watcher_extensions.py

@@ -0,0 +1,163 @@
+"""Tests for the bambu-watch extensions:
+
+- #9 temperature threshold reaches
+- #14 heartbeat (disconnect / reconnect) alarms
+"""
+
+from __future__ import annotations
+
+from unittest.mock import MagicMock
+
+import pytest
+
+from bambu_ai.watcher import HeartbeatState, ThresholdState
+
+# ---------------------------------------------------------------- ThresholdState
+
+
+class TestThresholdState:
+    def test_disabled_when_no_threshold_set(self) -> None:
+        state = ThresholdState()
+        notifier = MagicMock()
+        state.check_and_notify({"bed_temper": 500, "nozzle_temper": 500}, notifier)
+        notifier.assert_not_called()
+
+    def test_fires_once_when_bed_crosses_threshold(self) -> None:
+        state = ThresholdState(bed_threshold=65.0)
+        notifier = MagicMock()
+
+        # Below threshold — no fire.
+        state.check_and_notify({"bed_temper": 50.0}, notifier)
+        assert not state.bed_reached
+        notifier.assert_not_called()
+
+        # Crosses — one fire.
+        state.check_and_notify({"bed_temper": 65.5}, notifier)
+        assert state.bed_reached
+        assert notifier.call_count == 1
+        title = notifier.call_args.args[1]
+        assert "Bed" in title
+
+        # Stays above — no additional fires.
+        state.check_and_notify({"bed_temper": 65.0}, notifier)
+        state.check_and_notify({"bed_temper": 70.0}, notifier)
+        assert notifier.call_count == 1
+
+    def test_fires_once_when_nozzle_crosses_threshold(self) -> None:
+        state = ThresholdState(nozzle_threshold=220.0)
+        notifier = MagicMock()
+        state.check_and_notify({"nozzle_temper": 219.9}, notifier)
+        assert not state.nozzle_reached
+        state.check_and_notify({"nozzle_temper": 220.0}, notifier)
+        assert state.nozzle_reached
+        assert notifier.call_count == 1
+        assert "Nozzle" in notifier.call_args.args[1]
+
+    def test_bed_and_nozzle_are_independent(self) -> None:
+        state = ThresholdState(bed_threshold=65.0, nozzle_threshold=220.0)
+        notifier = MagicMock()
+        state.check_and_notify({"bed_temper": 65.0, "nozzle_temper": 0.0}, notifier)
+        state.check_and_notify({"bed_temper": 65.0, "nozzle_temper": 220.0}, notifier)
+        assert state.bed_reached and state.nozzle_reached
+        assert notifier.call_count == 2
+
+    def test_reset_for_new_print_re_arms(self) -> None:
+        state = ThresholdState(bed_threshold=65.0)
+        notifier = MagicMock()
+        state.check_and_notify({"bed_temper": 70.0}, notifier)
+        assert notifier.call_count == 1
+
+        state.reset_for_new_print()
+        assert not state.bed_reached
+        state.check_and_notify({"bed_temper": 70.0}, notifier)
+        assert notifier.call_count == 2
+
+    def test_ignores_non_numeric_values(self) -> None:
+        state = ThresholdState(bed_threshold=65.0)
+        notifier = MagicMock()
+        state.check_and_notify({"bed_temper": None}, notifier)
+        state.check_and_notify({"bed_temper": ""}, notifier)
+        state.check_and_notify({}, notifier)
+        notifier.assert_not_called()
+
+
+# ---------------------------------------------------------------- HeartbeatState
+
+
+class TestHeartbeatState:
+    def test_disabled_by_default_never_notifies(self) -> None:
+        state = HeartbeatState()
+        notifier = MagicMock()
+        state.on_disconnect(7, notifier)
+        state.on_connect(notifier)
+        notifier.assert_not_called()
+
+    def test_disconnect_with_nonzero_rc_fires_once_when_enabled(self) -> None:
+        state = HeartbeatState(enabled=True)
+        notifier = MagicMock()
+        state.on_disconnect(7, notifier)
+        assert state.in_alarm
+        assert notifier.call_count == 1
+        title = notifier.call_args.args[1]
+        assert "unreachable" in title.lower()
+
+    def test_repeated_disconnects_while_in_alarm_dont_re_notify(self) -> None:
+        state = HeartbeatState(enabled=True)
+        notifier = MagicMock()
+        state.on_disconnect(7, notifier)
+        state.on_disconnect(7, notifier)
+        state.on_disconnect(7, notifier)
+        assert notifier.call_count == 1
+        assert state.disconnect_count == 3
+
+    def test_clean_disconnect_rc_zero_does_not_notify(self) -> None:
+        """rc == 0 means we intentionally disconnected (Ctrl-C). No alert."""
+        state = HeartbeatState(enabled=True)
+        notifier = MagicMock()
+        state.on_disconnect(0, notifier)
+        notifier.assert_not_called()
+        assert not state.in_alarm
+
+    def test_reconnect_after_alarm_fires_back_online(self) -> None:
+        state = HeartbeatState(enabled=True)
+        notifier = MagicMock()
+        state.on_disconnect(7, notifier)
+        notifier.reset_mock()
+
+        state.on_connect(notifier)
+        assert not state.in_alarm
+        assert notifier.call_count == 1
+        title = notifier.call_args.args[1]
+        assert "back online" in title.lower()
+
+    def test_connect_without_prior_alarm_does_not_notify(self) -> None:
+        state = HeartbeatState(enabled=True)
+        notifier = MagicMock()
+        state.on_connect(notifier)
+        notifier.assert_not_called()
+
+
+# ---------------------------------------------------------------- arg parsing
+
+
+class TestArgParse:
+    def test_default_args_are_no_op(self) -> None:
+        from bambu_ai.watcher import _parse_args
+
+        a = _parse_args([])
+        assert a.notify_bed_temp is None
+        assert a.notify_nozzle_temp is None
+        assert a.heartbeat is False
+
+    def test_temperature_thresholds_parsed_as_float(self) -> None:
+        from bambu_ai.watcher import _parse_args
+
+        a = _parse_args(["--notify-bed-temp", "65", "--notify-nozzle-temp", "220.5"])
+        assert a.notify_bed_temp == pytest.approx(65.0)
+        assert a.notify_nozzle_temp == pytest.approx(220.5)
+
+    def test_heartbeat_flag(self) -> None:
+        from bambu_ai.watcher import _parse_args
+
+        a = _parse_args(["--heartbeat"])
+        assert a.heartbeat is True