diff --git a/tests/test_modbus_task.py b/tests/test_modbus_task.py
index 227cc71..5907ae7 100644
--- a/tests/test_modbus_task.py
+++ b/tests/test_modbus_task.py
@@ -1,11 +1,9 @@
 import pytest
-import threading
 from modpoll.arg_parser import get_parser
-from modpoll.modbus_task import ModbusMaster
+from modpoll.modbus_task import setup_modbus_handlers
 
 
 def test_modbus_task_modbus_setup():
-    modbus_masters = []
     parser = get_parser()
     args = parser.parse_args(
         [
@@ -16,14 +14,8 @@ def test_modbus_task_modbus_setup():
             "modsim.topmaker.net",
         ]
     )
-    event_exit = threading.Event()
-    for config_file in args.config:
-        modbus_master = ModbusMaster(args, event_exit, config_file)
-        if modbus_master.setup():
-            modbus_masters.append(modbus_master)
-        else:
-            modbus_master.close()
-    assert len(modbus_masters) == 2
+    modbus_handlers = setup_modbus_handlers(args)
+    assert len(modbus_handlers) == 2
 
 
 @pytest.mark.integration
@@ -37,22 +29,21 @@ def test_modbus_task_poll_modsim():
             "modsim.topmaker.net",
         ]
     )
-    event_exit = threading.Event()
-    modbus_master = ModbusMaster(args, event_exit, args.config[0])
+    modbus_handlers = setup_modbus_handlers(args)
+    modbus_handler = modbus_handlers[0]
 
-    assert modbus_master.setup()
-    assert modbus_master.connect()
+    assert modbus_handler.connect()
 
-    modbus_master.poll()
+    modbus_handler.poll()
 
     # Check if any data was collected
-    assert len(modbus_master.deviceList) > 0
-    assert len(modbus_master.deviceList[0].references) > 0
+    assert len(modbus_handler.deviceList) > 0
+    assert len(modbus_handler.deviceList[0].references) > 0
 
     # Check if at least one reference has a non-None value
     assert any(
-        ref.val is not None for ref in modbus_master.deviceList[0].references.values()
+        ref.val is not None for ref in modbus_handler.deviceList[0].references.values()
     )
 
-    modbus_master.disconnect()
-    modbus_master.close()
+    modbus_handler.disconnect()
+    modbus_handler.close()
diff --git a/tests/test_mqtt_task.py b/tests/test_mqtt_task.py
index 392e2b9..eab5759 100644
--- a/tests/test_mqtt_task.py
+++ b/tests/test_mqtt_task.py
@@ -4,7 +4,7 @@
 
 
 def test_mqtt_task_setup():
-    mqtt_task = MqttHandler(
+    mqtt_handler = MqttHandler(
         name="test_mqtt",
         host="mqtt.eclipseprojects.io",
         port=1883,
@@ -15,20 +15,20 @@ def test_mqtt_task_setup():
         subscribe_topics=["test/topic"],
     )
 
-    assert mqtt_task.mqttc_setup()
+    assert mqtt_handler.setup()
 
     # Test MQTT client properties
-    assert mqtt_task.mqttc is not None
-    assert mqtt_task.host == "mqtt.eclipseprojects.io"
-    assert mqtt_task.port == 1883
+    assert mqtt_handler.mqtt_client is not None
+    assert mqtt_handler.host == "mqtt.eclipseprojects.io"
+    assert mqtt_handler.port == 1883
 
     # Clean up
-    mqtt_task.mqttc_close()
+    mqtt_handler.close()
 
 
 @pytest.mark.integration
 def test_mqtt_task_connect():
-    mqtt_task = MqttHandler(
+    mqtt_handler = MqttHandler(
         name="test_mqtt",
         host="mqtt.eclipseprojects.io",
         port=1883,
@@ -38,14 +38,14 @@ def test_mqtt_task_connect():
         qos=0,
     )
 
-    assert mqtt_task.mqttc_setup()
-    assert mqtt_task.mqttc_connect()
+    assert mqtt_handler.setup()
+    assert mqtt_handler.connect()
 
     # Add a short delay to allow the connection to establish
     time.sleep(1)
 
     # Test connection
-    assert mqtt_task.mqttc_connected()
+    assert mqtt_handler.is_connected()
 
     # Clean up
-    mqtt_task.mqttc_close()
+    mqtt_handler.close()