diff --git a/src/bambu_ai/watcher.py b/src/bambu_ai/watcher.py index a4b1b39..567ede3 100644 --- a/src/bambu_ai/watcher.py +++ b/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,7 +47,136 @@ 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" @@ -47,6 +184,16 @@ def main() -> None: 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 diff --git a/tests/test_watcher_extensions.py b/tests/test_watcher_extensions.py new file mode 100644 index 0000000..74daccf --- /dev/null +++ b/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