From 5973b7ca26e63880e9d74a52fa9e2ec61c811401 Mon Sep 17 00:00:00 2001
From: Elena ``of Valhalla'' Grandi <valhalla@trueelena.org>
Date: Fri, 8 Dec 2023 10:20:36 +0100
Subject: Management command for the mqtt loop

---
 rrd/management/__init__.py             |  0
 rrd/management/commands/__init__.py    |  0
 rrd/management/commands/listen_mqtt.py | 10 ++++++
 rrd/mqtt.py                            | 64 ++++++++++++++++++++++++++++++++++
 rrd/tests.py                           |  3 --
 rrd/tests/__init__.py                  |  0
 rrd/tests/test_django.py               |  5 +++
 rrd/tests/test_mqtt.py                 | 30 ++++++++++++++++
 8 files changed, 109 insertions(+), 3 deletions(-)
 create mode 100644 rrd/management/__init__.py
 create mode 100644 rrd/management/commands/__init__.py
 create mode 100644 rrd/management/commands/listen_mqtt.py
 create mode 100644 rrd/mqtt.py
 delete mode 100644 rrd/tests.py
 create mode 100644 rrd/tests/__init__.py
 create mode 100644 rrd/tests/test_django.py
 create mode 100644 rrd/tests/test_mqtt.py

(limited to 'rrd')

diff --git a/rrd/management/__init__.py b/rrd/management/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/rrd/management/commands/__init__.py b/rrd/management/commands/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/rrd/management/commands/listen_mqtt.py b/rrd/management/commands/listen_mqtt.py
new file mode 100644
index 0000000..758a4f2
--- /dev/null
+++ b/rrd/management/commands/listen_mqtt.py
@@ -0,0 +1,10 @@
+from django.core.management.base import BaseCommand
+
+import rrd.mqtt
+
+class Command(BaseCommand):
+    help = "Listen to mqtt events and update rrd files and graphs"
+
+    def handle(self, *args, **kw):
+        mqtt_client = rrd.mqtt.MQTTClient()
+        mqtt_client.loop_forever()
diff --git a/rrd/mqtt.py b/rrd/mqtt.py
new file mode 100644
index 0000000..713d837
--- /dev/null
+++ b/rrd/mqtt.py
@@ -0,0 +1,64 @@
+import logging
+
+import django.conf
+import paho.mqtt.client as mqtt
+
+
+log = logging.getLogger(__name__)
+
+
+class MQTTClient:
+    def __init__(self):
+        self.reconnect = True
+        self.connected = False
+
+        self.client = mqtt.Client()
+        self.client.on_connect = self.on_connect
+        self.client.on_disconnect = self.on_disconnect
+        self.client.on_message = self.on_message
+
+    def loop_forever(self):
+        self.connect()
+        self.client.loop_forever()
+
+    def loop_start(self):
+        self.connect()
+        self.client.loop_start()
+
+    def loop_stop(self):
+        self.disconnect(reconnect=False)
+        self.client.loop_stop()
+
+    def connect(self):
+        try:
+            self.client.connect(
+                django.conf.settings.MQTT_SERVER,
+                django.conf.settings.MQTT_PORT,
+                60,  # TODO: make the keepalive configurable
+            )
+        except OSError as e:
+            log.debug("Could not connect to MQTT server")
+            log.debug(e)
+
+    def disconnect(self, reconnect: bool = True):
+        self.reconnect = reconnect
+        self.client.disconnect()
+
+    def on_connect(self, client, userdata, flags, rc):
+        log.info("Connected to MQTT")
+        self.connected = True
+        client.subscribe(django.conf.settings.MQTT_TOPIC)
+
+    def on_disconnect(self, client, userdata, rc):
+        log.info("Disconnected from MQTT")
+        self.connected = False
+        if self.reconnect:
+            self.connect()
+
+    def on_message(self, client, userdata, msg):
+        print(
+            "Received msg %s %s", msg.topic, msg.payload.decode()
+        )
+        log.debug(
+            "Received msg %s %s", msg.topic, msg.payload.decode()
+        )
diff --git a/rrd/tests.py b/rrd/tests.py
deleted file mode 100644
index 7ce503c..0000000
--- a/rrd/tests.py
+++ /dev/null
@@ -1,3 +0,0 @@
-from django.test import TestCase
-
-# Create your tests here.
diff --git a/rrd/tests/__init__.py b/rrd/tests/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/rrd/tests/test_django.py b/rrd/tests/test_django.py
new file mode 100644
index 0000000..e0fbb45
--- /dev/null
+++ b/rrd/tests/test_django.py
@@ -0,0 +1,5 @@
+from django.test import TestCase
+
+# Create your tests here.
+
+
diff --git a/rrd/tests/test_mqtt.py b/rrd/tests/test_mqtt.py
new file mode 100644
index 0000000..4ea0852
--- /dev/null
+++ b/rrd/tests/test_mqtt.py
@@ -0,0 +1,30 @@
+import time
+
+import django.test
+
+from .. import mqtt
+
+class TestMQTT(django.test.TestCase):
+    def setUp(self):
+        self.mqtt = mqtt.MQTTClient()
+        self.mqtt.loop_start()
+        time.sleep(0.1)
+        if not self.mqtt.connected:
+            self.skipTest("Could not find an mqtt server")
+
+    def tearDown(self):
+        self.mqtt.loop_stop()
+
+    def test_disconnect(self):
+        # after disconnecting from the mqtt server, we should
+        # automatically reconnect
+        self.mqtt.disconnect()
+        time.sleep(2)
+        self.assertTrue(self.mqtt.connected)
+
+    def test_disconnect_and_stay(self):
+        # unless we really want to force a disconnection
+        self.mqtt.disconnect(reconnect=False)
+        time.sleep(2)
+        self.assertFalse(self.mqtt.connected)
+        self.assertFalse(self.mqtt.reconnect)
-- 
cgit v1.2.3