From 0a5edb8319de340c4b8bafde53a39e2f237b588a Mon Sep 17 00:00:00 2001 From: Ruairi Carroll Date: Tue, 7 Sep 2021 20:56:30 +0100 Subject: [PATCH 1/2] - Use logging instead of print() - Force fetch_from_config to raise ValueError on missing keys in config. - Fixup tests --- wgkex/worker/mqtt.py | 24 ++++++++++++++++++------ wgkex/worker/mqtt_test.py | 11 +++-------- 2 files changed, 21 insertions(+), 14 deletions(-) diff --git a/wgkex/worker/mqtt.py b/wgkex/worker/mqtt.py index ee3d8cc..5176c08 100644 --- a/wgkex/worker/mqtt.py +++ b/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)) diff --git a/wgkex/worker/mqtt_test.py b/wgkex/worker/mqtt_test.py index e5c1b0d..85ffdf9 100644 --- a/wgkex/worker/mqtt_test.py +++ b/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, ) From e49b501d44d62b9bf838a994e61eec35eea605ff Mon Sep 17 00:00:00 2001 From: Annika Wickert Date: Tue, 7 Sep 2021 22:05:15 +0200 Subject: [PATCH 2/2] Update wgkex/worker/mqtt.py --- wgkex/worker/mqtt.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/wgkex/worker/mqtt.py b/wgkex/worker/mqtt.py index 5176c08..5d441ec 100644 --- a/wgkex/worker/mqtt.py +++ b/wgkex/worker/mqtt.py @@ -75,4 +75,4 @@ def on_message(client: mqtt.Client, userdata: Any, message: mqtt.MQTTMessage) -> ) 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))