From 485bc4468ee3c1bfb07058be8e254eb134f03580 Mon Sep 17 00:00:00 2001
From: Maaike <maaike.limper@gmail.com>
Date: Thu, 25 Jul 2024 17:51:26 +0200
Subject: [PATCH] add environment test to catch config issues

---
 wis2box-management/docker/entrypoint.sh   |  3 +++
 wis2box-management/wis2box/env.py         | 32 ++++++++++++++++++++++
 wis2box-management/wis2box/pubsub/base.py | 12 +++++++++
 wis2box-management/wis2box/pubsub/mqtt.py | 33 +++++++++++++++++++++++
 4 files changed, 80 insertions(+)

diff --git a/wis2box-management/docker/entrypoint.sh b/wis2box-management/docker/entrypoint.sh
index 8baf4c223..4766c8bd2 100755
--- a/wis2box-management/docker/entrypoint.sh
+++ b/wis2box-management/docker/entrypoint.sh
@@ -38,6 +38,9 @@ wis2box environment create
 wis2box environment show
 wis2box api setup
 
+# test the wis2box is not misconfigured
+wis2box environment test
+
 # ensure cron is running
 service cron start
 service cron status
diff --git a/wis2box-management/wis2box/env.py b/wis2box-management/wis2box/env.py
index c8778c6d2..54646d336 100644
--- a/wis2box-management/wis2box/env.py
+++ b/wis2box-management/wis2box/env.py
@@ -108,6 +108,37 @@ def environment():
     pass
 
 
+@click.command()
+@click.pass_context
+@cli_helpers.OPTION_VERBOSITY
+def test(ctx, verbosity):
+    """Tests the environment is set up correctly"""
+
+    click.echo(f'Setting up logging (loglevel={LOGLEVEL}, logfile={LOGFILE})')
+    setup_logger(LOGLEVEL, LOGFILE)
+
+    click.echo('Testing BROKER_PUBLIC')
+    # load plugin for plugin-broker
+    defs = {
+        'codepath': PLUGINS['pubsub']['mqtt']['plugin'],
+        'url': BROKER_PUBLIC,
+        'client_type': 'publisher'
+    }
+    broker = load_plugin('pubsub', defs)
+
+    try:
+        result = broker.test(topic='origin/a/wis2/test', message='wis2box pub test')
+    except Exception as err:
+        LOGGER.error(err)
+        raise EnvironmentError(err)
+
+    if result:
+        click.echo('Broker test successful')
+    else:
+        LOGGER.error('Could not connect to broker defined by WI2BOX_BROKER_PUBLIC')
+        click.echo('Broker test failed')
+        exit(1)
+
 @click.command()
 @click.pass_context
 @cli_helpers.OPTION_VERBOSITY
@@ -163,3 +194,4 @@ def show(ctx, verbosity):
 
 environment.add_command(create)
 environment.add_command(show)
+environment.add_command(test)
diff --git a/wis2box-management/wis2box/pubsub/base.py b/wis2box-management/wis2box/pubsub/base.py
index 6f92187b1..29edea4ec 100644
--- a/wis2box-management/wis2box/pubsub/base.py
+++ b/wis2box-management/wis2box/pubsub/base.py
@@ -64,6 +64,18 @@ def sub(self, topic: str) -> None:
 
         raise NotImplementedError()
 
+    def test(self, topic='wis2box/test', message='test') -> bool:
+        """
+        Test the connection to the broker
+
+        :param topic: `str` of topic
+        :param message: `str` of message
+
+        :returns: `bool` of test result
+        """
+
+        raise NotImplementedError()
+
     def bind(self, event: str, function: Callable[..., Any]) -> None:
         """
         Binds an event to a function
diff --git a/wis2box-management/wis2box/pubsub/mqtt.py b/wis2box-management/wis2box/pubsub/mqtt.py
index 3abb7a0d8..5e8d1103a 100644
--- a/wis2box-management/wis2box/pubsub/mqtt.py
+++ b/wis2box-management/wis2box/pubsub/mqtt.py
@@ -22,6 +22,7 @@
 import logging
 import random
 
+from time import sleep
 from typing import Any, Callable
 
 from paho.mqtt import client as mqtt_client
@@ -43,6 +44,7 @@ def __init__(self, broker: str) -> None:
         """
 
         super().__init__(broker)
+        self.test_status = 'unknown'
         self.type = 'mqtt'
         self._port = self.broker_url.port
         self.client_id = f"wis2box-mqtt-{self.broker['client_type']}-{random.randint(0, 1000)}"  # noqa
@@ -136,5 +138,36 @@ def bind(self, event: str, function: Callable[..., Any]) -> None:
 
         setattr(self.conn, event, function)
 
+    def test(self, topic='wis2box/test', message='test') -> bool:
+        """
+        Test the connection to the broker
+
+        :returns: `bool` of test result
+        """
+
+        def on_connect(client, userdata, flags, rc):
+            if rc == 0:
+                LOGGER.debug(f'Test: Connected to broker {self.broker}')
+                client.subscribe(topic, qos=1)
+                LOGGER.debug(f'Test: Subscribed to topic {topic}')
+            else:
+                msg = f'Test: Failed to connect to MQTT-broker: {mqtt_client.connack_string(rc)}' # noqa
+                LOGGER.error(msg)
+        
+        def on_message(client, userdata, message):
+            LOGGER.debug(f'Test: Received message {message.payload.decode()}')
+            self.test_status = 'success'
+
+        self.conn.on_connect = on_connect
+        self.conn.on_message = on_message
+    
+        self.conn.loop_start()
+        sleep(0.1)
+        self.conn.publish(topic, message, qos=1)
+        sleep(0.1)
+        self.conn.loop_stop()
+
+        return self.test_status == 'success'
+
     def __repr__(self):
         return '<MQTTPubSubClient>'