diff --git a/demo.py b/demo.py
index e29df29..a644e73 100755
--- a/demo.py
+++ b/demo.py
@@ -161,7 +161,6 @@ def get_device_state(device_info, args):
     return state
 
 
-# pylint: disable=too-many-statements,too-many-locals
 def main():
     """
     Main loop. Acquire state from plugs using MQTT and use a threshold to determine
@@ -176,7 +175,26 @@ def main():
     # map of device name to Devices instance
     devices = {}
 
-    # connect to MQTT broker
+    mqtt = mqtt_setup(args, devices)
+
+    i = 0
+    while True:
+        # Make sure to stay connected to the broker e.g. in case of keep alive.
+        i += 1
+        logger.debug(f"Loop {i}")
+        mqtt.loop(1)
+
+        for device_name, device_info in devices.items():
+            logger.debug(f"{device_info}")
+            logger.info(f"{device_name} = {get_device_state(device_info, args)}")
+
+
+def mqtt_setup(args, devices):
+    """
+    connect to MQTT broker and subscribe to topic given by the args
+    """
+    logger = logging.getLogger(__name__)
+
     mqtt = MQTT.MQTT(
         broker=args.hostname,
         port=args.port,
@@ -184,25 +202,16 @@ def main():
         ssl_context=ssl.create_default_context(),
         user_data=devices,
     )
-
     mqtt.on_connect = connect
     mqtt.on_subscribe = subscribe
     mqtt.on_message = message
 
     logger.info(f"Connecting to MQTT broker {args.hostname} on port {args.port}")
     mqtt.connect()
-    mqtt.subscribe(args.topic, qos=0)
 
-    i = 0
-    while True:
-        # Make sure to stay connected to the broker e.g. in case of keep alive.
-        i += 1
-        logger.debug(f"Loop {i}")
-        mqtt.loop(1)
+    mqtt.subscribe(args.topic, qos=0)
 
-        for device_name, device_info in devices.items():
-            logger.debug(f"{device_info}")
-            logger.info(f"{device_name} = {get_device_state(device_info, args)}")
+    return mqtt
 
 
 if __name__ == "__main__":