diff --git a/tests/connect_publish.c b/tests/connect_publish.c
index 4d2d43b..b91c341 100644
--- a/tests/connect_publish.c
+++ b/tests/connect_publish.c
@@ -4,6 +4,8 @@
 
 #include "mqtt.h"
 
+bool leave = false;
+
 #define LOG(fmt, ...) fprintf(stdout, fmt "\n", ## __VA_ARGS__)
 
 bool err_handler(MQTTHandle *handle, MQTTErrorCode error) {
@@ -13,6 +15,28 @@ bool err_handler(MQTTHandle *handle, MQTTErrorCode error) {
     return true;
 }
 
+void mqtt_connected(MQTTHandle *handle, void *context) {
+    LOG("Connected!");
+
+    LOG("Trying publish to testsuite/mqtt/test...");
+    MQTTStatus result = mqtt_publish(handle, "testsuite/mqtt/test", "payload", MQTT_QOS_0);
+    if (result != MQTT_STATUS_OK) {
+        LOG("Could not publish");
+        exit(1);
+    }
+
+    sleep(1);
+
+    LOG("Disconnecting...");
+    result = mqtt_disconnect(handle, NULL, NULL);
+    if (result != MQTT_STATUS_OK) {
+        LOG("Could not disconnect");
+        exit(1);
+    }
+
+    exit(0);
+}
+
 int main(int argc, char **argv) {
     MQTTConfig config = { 0 };
 
@@ -21,31 +45,15 @@ int main(int argc, char **argv) {
     config.port = 1883;
 
     LOG("Trying to connect to %s", config.hostname);
-    MQTTHandle *mqtt = mqtt_connect(&config, err_handler);
+    MQTTHandle *mqtt = mqtt_connect(&config, mqtt_connected, NULL, err_handler);
 
     if (mqtt == NULL) {
         LOG("Connection failed!");
         return 1;
     }
 
-    LOG("Connected!");
-
-    sleep(5);
-
-    LOG("Trying publish to testsuite/mqtt/test...");
-    MQTTStatus result = mqtt_publish(mqtt, "testsuite/mqtt/test", "payload", MQTT_QOS_0);
-    if (result != MQTT_STATUS_OK) {
-        LOG("Could not publish");
-        return 1;
+    while (!leave) {
+        LOG("Waiting...");
+        sleep(1);
     }
-
-    sleep(5);
-
-    LOG("Disconnecting...");
-    result = mqtt_disconnect(mqtt);
-    if (result != MQTT_STATUS_OK) {
-        LOG("Could not disconnect");
-        return 1;
-    }
-    return 0;
 }
diff --git a/tests/connect_subscribe.c b/tests/connect_subscribe.c
index 2c87cc1..ad17f0b 100644
--- a/tests/connect_subscribe.c
+++ b/tests/connect_subscribe.c
@@ -1,46 +1,49 @@
 #include <stdio.h>
+#include <stdlib.h>
 #include <unistd.h>
 
 #include "mqtt.h"
 
+bool leave = false;
+
 #define LOG(fmt, ...) fprintf(stdout, fmt "\n", ## __VA_ARGS__)
 
 bool err_handler(MQTTHandle *handle, MQTTErrorCode error) {
     LOG("Error received: %d", error);
-    
+
     return 1;
 }
 
-bool leave = false;
 void callback(MQTTHandle *handle, char *topic, char *payload) {
     LOG("Received publish: %s -> %s", topic, payload);
     leave = true;
 }
 
+void mqtt_connected(MQTTHandle *handle, void *context) {
+    LOG("Connected!");
+
+    LOG("Trying subscribe on testsuite/mqtt/test...");
+    MQTTStatus result = mqtt_subscribe(handle, "testsuite/mqtt/test", callback);
+    if (result != MQTT_STATUS_OK) {
+        LOG("Could not subscribe");
+        exit(1);
+    }
+}
+
 int main(int argc, char **argv) {
-    MQTTConfig config;
+    MQTTConfig config = { 0 };
 
     config.client_id = "libmqtt_testsuite";
-    config.hostname = "test.mosquitto.org";
+    config.hostname = "localhost";
     config.port = 1883;
 
-    LOG("Trying to connect to test.mosquitto.org");
-    MQTTHandle *mqtt = mqtt_connect(&config, err_handler);
+    LOG("Trying to connect to %s...", config.hostname);
+    MQTTHandle *mqtt = mqtt_connect(&config, mqtt_connected, NULL, err_handler);
 
     if (mqtt == NULL) {
         LOG("Connection failed!");
         return 1;
     }
-    LOG("Connected!");
-    
-    sleep(5);
-
-    LOG("Trying subscribe on testsuite/mqtt/test...");
-    MQTTStatus result = mqtt_subscribe(mqtt, "testsuite/mqtt/test", callback);
-    if (result != MQTT_STATUS_OK) {
-        LOG("Could not publish");
-        return 1;
-    }
 
     while (!leave) {
         LOG("Waiting...");
@@ -48,6 +51,5 @@ int main(int argc, char **argv) {
     }
 
     LOG("Disconnecting...");
-    mqtt_disconnect(mqtt);
-    return 0;
+    mqtt_disconnect(mqtt, NULL, NULL);
 }