reverted TransportSqlTestSuite typo from previous commit & updated provision tests

application/src/test/java/org/thingsboard/server/transport/TransportSqlTestSuite.java

@@ -20,15 +20,15 @@ import org.junit.runner.RunWith;
 
 @RunWith(ClasspathSuite.class)
 @ClasspathSuite.ClassnameFilters({
-//        "org.thingsboard.server.transport.*.rpc.*Test",
+        "org.thingsboard.server.transport.*.rpc.*Test",
-//        "org.thingsboard.server.transport.*.telemetry.timeseries.sql.*Test",
+        "org.thingsboard.server.transport.*.telemetry.timeseries.sql.*Test",
-//        "org.thingsboard.server.transport.*.telemetry.attributes.*Test",
+        "org.thingsboard.server.transport.*.telemetry.attributes.*Test",
-//        "org.thingsboard.server.transport.*.attributes.updates.*Test",
+        "org.thingsboard.server.transport.*.attributes.updates.*Test",
-//        "org.thingsboard.server.transport.*.attributes.request.*Test",
+        "org.thingsboard.server.transport.*.attributes.request.*Test",
-        "org.thingsboard.server.transport.mqtt.claim.*Test",
+        "org.thingsboard.server.transport.*.claim.*Test",
-//        "org.thingsboard.server.transport.*.provision.*Test",
+        "org.thingsboard.server.transport.*.provision.*Test",
-//        "org.thingsboard.server.transport.*.credentials.*Test",
+        "org.thingsboard.server.transport.*.credentials.*Test",
-//        "org.thingsboard.server.transport.lwm2m.*.sql.*Test"
+        "org.thingsboard.server.transport.lwm2m.*.sql.*Test"
 })
 public class TransportSqlTestSuite {
 

application/src/test/java/org/thingsboard/server/transport/mqtt/provision/MqttProvisionJsonDeviceTest.java

@@ -15,13 +15,9 @@
  */
 package org.thingsboard.server.transport.mqtt.provision;
 
-import com.google.gson.JsonObject;
+import com.fasterxml.jackson.databind.JsonNode;
 import io.netty.handler.codec.mqtt.MqttQoS;
 import lombok.extern.slf4j.Slf4j;
-import org.eclipse.paho.client.mqttv3.IMqttDeliveryToken;
-import org.eclipse.paho.client.mqttv3.MqttAsyncClient;
-import org.eclipse.paho.client.mqttv3.MqttCallback;
-import org.eclipse.paho.client.mqttv3.MqttMessage;
 import org.junit.Assert;
 import org.junit.Test;
 import org.springframework.beans.factory.annotation.Autowired;
@@ -30,20 +26,22 @@ import org.thingsboard.server.common.data.Device;
 import org.thingsboard.server.common.data.DeviceProfileProvisionType;
 import org.thingsboard.server.common.data.TransportPayloadType;
 import org.thingsboard.server.common.data.device.credentials.BasicMqttCredentials;
-import org.thingsboard.server.common.data.device.profile.MqttTopics;
 import org.thingsboard.server.common.data.security.DeviceCredentials;
 import org.thingsboard.server.common.msg.EncryptionUtil;
-import org.thingsboard.server.common.transport.util.JsonUtils;
 import org.thingsboard.server.dao.device.DeviceCredentialsService;
 import org.thingsboard.server.dao.device.DeviceService;
 import org.thingsboard.server.dao.device.provision.ProvisionResponseStatus;
 import org.thingsboard.server.dao.service.DaoSqlTest;
 import org.thingsboard.server.transport.mqtt.AbstractMqttIntegrationTest;
+import org.thingsboard.server.transport.mqtt.MqttTestCallback;
+import org.thingsboard.server.transport.mqtt.MqttTestClient;
 import org.thingsboard.server.transport.mqtt.MqttTestConfigProperties;
 
-import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.TimeUnit;
 
+import static org.thingsboard.server.common.data.device.profile.MqttTopics.DEVICE_PROVISION_REQUEST_TOPIC;
+import static org.thingsboard.server.common.data.device.profile.MqttTopics.DEVICE_PROVISION_RESPONSE_TOPIC;
+
 @Slf4j
 @DaoSqlTest
 public class MqttProvisionJsonDeviceTest extends AbstractMqttIntegrationTest {
@@ -97,10 +95,12 @@ public class MqttProvisionJsonDeviceTest extends AbstractMqttIntegrationTest {
                 .provisionType(DeviceProfileProvisionType.DISABLED)
                 .build();
         super.processBeforeTest(configProperties);
-        byte[] result = createMqttClientAndPublish().getPayloadBytes();
+        byte[] result = createMqttClientAndPublish();
-        JsonObject response = JsonUtils.parse(new String(result)).getAsJsonObject();
+        JsonNode response = JacksonUtil.fromBytes(result);
-        Assert.assertEquals("Provision data was not found!", response.get("errorMsg").getAsString());
+        Assert.assertTrue(response.hasNonNull("errorMsg"));
-        Assert.assertEquals(ProvisionResponseStatus.NOT_FOUND.name(), response.get("status").getAsString());
+        Assert.assertTrue(response.hasNonNull("status"));
+        Assert.assertEquals("Provision data was not found!", response.get("errorMsg").asText());
+        Assert.assertEquals(ProvisionResponseStatus.NOT_FOUND.name(), response.get("status").asText());
     }
 
 
@@ -113,8 +113,10 @@ public class MqttProvisionJsonDeviceTest extends AbstractMqttIntegrationTest {
                 .provisionSecret("testProvisionSecret")
                 .build();
         super.processBeforeTest(configProperties);
-        byte[] result = createMqttClientAndPublish().getPayloadBytes();
+        byte[] result = createMqttClientAndPublish();
-        JsonObject response = JsonUtils.parse(new String(result)).getAsJsonObject();
+        JsonNode response = JacksonUtil.fromBytes(result);
+        Assert.assertTrue(response.hasNonNull("credentialsType"));
+        Assert.assertTrue(response.hasNonNull("status"));
 
         Device createdDevice = deviceService.findDeviceByTenantIdAndName(tenantId, "Test Provision device");
 
@@ -122,8 +124,8 @@ public class MqttProvisionJsonDeviceTest extends AbstractMqttIntegrationTest {
 
         DeviceCredentials deviceCredentials = deviceCredentialsService.findDeviceCredentialsByDeviceId(tenantId, createdDevice.getId());
 
-        Assert.assertEquals(deviceCredentials.getCredentialsType().name(), response.get("credentialsType").getAsString());
+        Assert.assertEquals(deviceCredentials.getCredentialsType().name(), response.get("credentialsType").asText());
-        Assert.assertEquals(ProvisionResponseStatus.SUCCESS.name(), response.get("status").getAsString());
+        Assert.assertEquals(ProvisionResponseStatus.SUCCESS.name(), response.get("status").asText());
     }
 
 
@@ -137,8 +139,10 @@ public class MqttProvisionJsonDeviceTest extends AbstractMqttIntegrationTest {
                 .build();
         super.processBeforeTest(configProperties);
         String requestCredentials = ",\"credentialsType\": \"ACCESS_TOKEN\",\"token\": \"test_token\"";
-        byte[] result = createMqttClientAndPublish(requestCredentials).getPayloadBytes();
+        byte[] result = createMqttClientAndPublish(requestCredentials);
-        JsonObject response = JsonUtils.parse(new String(result)).getAsJsonObject();
+        JsonNode response = JacksonUtil.fromBytes(result);
+        Assert.assertTrue(response.hasNonNull("credentialsType"));
+        Assert.assertTrue(response.hasNonNull("status"));
 
         Device createdDevice = deviceService.findDeviceByTenantIdAndName(tenantId, "Test Provision device");
 
@@ -146,10 +150,10 @@ public class MqttProvisionJsonDeviceTest extends AbstractMqttIntegrationTest {
 
         DeviceCredentials deviceCredentials = deviceCredentialsService.findDeviceCredentialsByDeviceId(tenantId, createdDevice.getId());
 
-        Assert.assertEquals(deviceCredentials.getCredentialsType().name(), response.get("credentialsType").getAsString());
+        Assert.assertEquals(deviceCredentials.getCredentialsType().name(), response.get("credentialsType").asText());
         Assert.assertEquals(deviceCredentials.getCredentialsType().name(), "ACCESS_TOKEN");
         Assert.assertEquals(deviceCredentials.getCredentialsId(), "test_token");
-        Assert.assertEquals(ProvisionResponseStatus.SUCCESS.name(), response.get("status").getAsString());
+        Assert.assertEquals(ProvisionResponseStatus.SUCCESS.name(), response.get("status").asText());
     }
 
 
@@ -163,8 +167,10 @@ public class MqttProvisionJsonDeviceTest extends AbstractMqttIntegrationTest {
                 .build();
         super.processBeforeTest(configProperties);
         String requestCredentials = ",\"credentialsType\": \"X509_CERTIFICATE\",\"hash\": \"testHash\"";
-        byte[] result = createMqttClientAndPublish(requestCredentials).getPayloadBytes();
+        byte[] result = createMqttClientAndPublish(requestCredentials);
-        JsonObject response = JsonUtils.parse(new String(result)).getAsJsonObject();
+        JsonNode response = JacksonUtil.fromBytes(result);
+        Assert.assertTrue(response.hasNonNull("credentialsType"));
+        Assert.assertTrue(response.hasNonNull("status"));
 
         Device createdDevice = deviceService.findDeviceByTenantIdAndName(tenantId, "Test Provision device");
 
@@ -172,7 +178,7 @@ public class MqttProvisionJsonDeviceTest extends AbstractMqttIntegrationTest {
 
         DeviceCredentials deviceCredentials = deviceCredentialsService.findDeviceCredentialsByDeviceId(tenantId, createdDevice.getId());
 
-        Assert.assertEquals(deviceCredentials.getCredentialsType().name(), response.get("credentialsType").getAsString());
+        Assert.assertEquals(deviceCredentials.getCredentialsType().name(), response.get("credentialsType").asText());
         Assert.assertEquals(deviceCredentials.getCredentialsType().name(), "X509_CERTIFICATE");
 
         String cert = EncryptionUtil.certTrimNewLines(deviceCredentials.getCredentialsValue());
@@ -181,7 +187,7 @@ public class MqttProvisionJsonDeviceTest extends AbstractMqttIntegrationTest {
         Assert.assertEquals(deviceCredentials.getCredentialsId(), sha3Hash);
 
         Assert.assertEquals(deviceCredentials.getCredentialsValue(), "testHash");
-        Assert.assertEquals(ProvisionResponseStatus.SUCCESS.name(), response.get("status").getAsString());
+        Assert.assertEquals(ProvisionResponseStatus.SUCCESS.name(), response.get("status").asText());
     }
 
 
@@ -195,8 +201,10 @@ public class MqttProvisionJsonDeviceTest extends AbstractMqttIntegrationTest {
                 .build();
         super.processBeforeTest(configProperties);
         String requestCredentials = ",\"credentialsType\": \"MQTT_BASIC\",\"clientId\": \"test_clientId\",\"username\": \"test_username\",\"password\": \"test_password\"";
-        byte[] result = createMqttClientAndPublish(requestCredentials).getPayloadBytes();
+        byte[] result = createMqttClientAndPublish(requestCredentials);
-        JsonObject response = JsonUtils.parse(new String(result)).getAsJsonObject();
+        JsonNode response = JacksonUtil.fromBytes(result);
+        Assert.assertTrue(response.hasNonNull("credentialsType"));
+        Assert.assertTrue(response.hasNonNull("status"));
 
         Device createdDevice = deviceService.findDeviceByTenantIdAndName(tenantId, "Test Provision device");
 
@@ -204,7 +212,7 @@ public class MqttProvisionJsonDeviceTest extends AbstractMqttIntegrationTest {
 
         DeviceCredentials deviceCredentials = deviceCredentialsService.findDeviceCredentialsByDeviceId(tenantId, createdDevice.getId());
 
-        Assert.assertEquals(deviceCredentials.getCredentialsType().name(), response.get("credentialsType").getAsString());
+        Assert.assertEquals(deviceCredentials.getCredentialsType().name(), response.get("credentialsType").asText());
         Assert.assertEquals(deviceCredentials.getCredentialsType().name(), "MQTT_BASIC");
         Assert.assertEquals(deviceCredentials.getCredentialsId(), EncryptionUtil.getSha3Hash("|", "test_clientId", "test_username"));
 
@@ -214,7 +222,7 @@ public class MqttProvisionJsonDeviceTest extends AbstractMqttIntegrationTest {
         mqttCredentials.setPassword("test_password");
 
         Assert.assertEquals(deviceCredentials.getCredentialsValue(), JacksonUtil.toString(mqttCredentials));
-        Assert.assertEquals(ProvisionResponseStatus.SUCCESS.name(), response.get("status").getAsString());
+        Assert.assertEquals(ProvisionResponseStatus.SUCCESS.name(), response.get("status").asText());
     }
 
     protected void processTestProvisioningCheckPreProvisionedDevice() throws Exception {
@@ -226,13 +234,15 @@ public class MqttProvisionJsonDeviceTest extends AbstractMqttIntegrationTest {
                 .provisionSecret("testProvisionSecret")
                 .build();
         super.processBeforeTest(configProperties);
-        byte[] result = createMqttClientAndPublish().getPayloadBytes();
+        byte[] result = createMqttClientAndPublish();
-        JsonObject response = JsonUtils.parse(new String(result)).getAsJsonObject();
+        JsonNode response = JacksonUtil.fromBytes(result);
+        Assert.assertTrue(response.hasNonNull("credentialsType"));
+        Assert.assertTrue(response.hasNonNull("status"));
 
         DeviceCredentials deviceCredentials = deviceCredentialsService.findDeviceCredentialsByDeviceId(tenantId, savedDevice.getId());
 
-        Assert.assertEquals(deviceCredentials.getCredentialsType().name(), response.get("credentialsType").getAsString());
+        Assert.assertEquals(deviceCredentials.getCredentialsType().name(), response.get("credentialsType").asText());
-        Assert.assertEquals(ProvisionResponseStatus.SUCCESS.name(), response.get("status").getAsString());
+        Assert.assertEquals(ProvisionResponseStatus.SUCCESS.name(), response.get("status").asText());
     }
 
     protected void processTestProvisioningWithBadKeyDevice() throws Exception {
@@ -244,72 +254,29 @@ public class MqttProvisionJsonDeviceTest extends AbstractMqttIntegrationTest {
                 .provisionSecret("testProvisionSecret")
                 .build();
         super.processBeforeTest(configProperties);
-        byte[] result = createMqttClientAndPublish().getPayloadBytes();
+        byte[] result = createMqttClientAndPublish();
-        JsonObject response = JsonUtils.parse(new String(result)).getAsJsonObject();
+        JsonNode response = JacksonUtil.fromBytes(result);
-        Assert.assertEquals("Provision data was not found!", response.get("errorMsg").getAsString());
+        Assert.assertTrue(response.hasNonNull("errorMsg"));
-        Assert.assertEquals(ProvisionResponseStatus.NOT_FOUND.name(), response.get("status").getAsString());
+        Assert.assertTrue(response.hasNonNull("status"));
+        Assert.assertEquals("Provision data was not found!", response.get("errorMsg").asText());
+        Assert.assertEquals(ProvisionResponseStatus.NOT_FOUND.name(), response.get("status").asText());
     }
 
-    protected TestMqttCallback createMqttClientAndPublish() throws Exception {
+    protected byte[] createMqttClientAndPublish() throws Exception {
         return createMqttClientAndPublish("");
     }
 
-    protected TestMqttCallback createMqttClientAndPublish(String deviceCredentials) throws Exception {
+    protected byte[] createMqttClientAndPublish(String deviceCredentials) throws Exception {
         String provisionRequestMsg = createTestProvisionMessage(deviceCredentials);
-        MqttAsyncClient client = getMqttAsyncClient("provision");
+        MqttTestClient client = new MqttTestClient();
-        TestMqttCallback onProvisionCallback = getTestMqttCallback();
+        client.connectAndWait("provision");
+        MqttTestCallback onProvisionCallback = new MqttTestCallback(DEVICE_PROVISION_RESPONSE_TOPIC);
         client.setCallback(onProvisionCallback);
-        client.subscribe(MqttTopics.DEVICE_PROVISION_RESPONSE_TOPIC, MqttQoS.AT_MOST_ONCE.value());
+        client.subscribe(DEVICE_PROVISION_RESPONSE_TOPIC, MqttQoS.AT_MOST_ONCE);
-        Thread.sleep(2000);
+        client.publishAndWait(DEVICE_PROVISION_REQUEST_TOPIC, provisionRequestMsg.getBytes());
-        client.publish(MqttTopics.DEVICE_PROVISION_REQUEST_TOPIC, new MqttMessage(provisionRequestMsg.getBytes()));
+        onProvisionCallback.getSubscribeLatch().await(3, TimeUnit.SECONDS);
-        onProvisionCallback.getLatch().await(3, TimeUnit.SECONDS);
+        client.disconnect();
-        return onProvisionCallback;
+        return onProvisionCallback.getPayloadBytes();
-    }
-
-
-    protected TestMqttCallback getTestMqttCallback() {
-        CountDownLatch latch = new CountDownLatch(1);
-        return new TestMqttCallback(latch);
-    }
-
-
-    protected static class TestMqttCallback implements MqttCallback {
-
-        private final CountDownLatch latch;
-        private Integer qoS;
-        private byte[] payloadBytes;
-
-        TestMqttCallback(CountDownLatch latch) {
-            this.latch = latch;
-        }
-
-        public int getQoS() {
-            return qoS;
-        }
-
-        public byte[] getPayloadBytes() {
-            return payloadBytes;
-        }
-
-        public CountDownLatch getLatch() {
-            return latch;
-        }
-
-        @Override
-        public void connectionLost(Throwable throwable) {
-        }
-
-        @Override
-        public void messageArrived(String requestTopic, MqttMessage mqttMessage) throws Exception {
-            qoS = mqttMessage.getQos();
-            payloadBytes = mqttMessage.getPayload();
-            latch.countDown();
-        }
-
-        @Override
-        public void deliveryComplete(IMqttDeliveryToken iMqttDeliveryToken) {
-
-        }
     }
 
     protected String createTestProvisionMessage(String deviceCredentials) {

application/src/test/java/org/thingsboard/server/transport/mqtt/provision/MqttProvisionProtoDeviceTest.java

@@ -17,10 +17,6 @@ package org.thingsboard.server.transport.mqtt.provision;
 
 import io.netty.handler.codec.mqtt.MqttQoS;
 import lombok.extern.slf4j.Slf4j;
-import org.eclipse.paho.client.mqttv3.IMqttDeliveryToken;
-import org.eclipse.paho.client.mqttv3.MqttAsyncClient;
-import org.eclipse.paho.client.mqttv3.MqttCallback;
-import org.eclipse.paho.client.mqttv3.MqttMessage;
 import org.junit.Assert;
 import org.junit.Test;
 import org.springframework.beans.factory.annotation.Autowired;
@@ -29,7 +25,6 @@ import org.thingsboard.server.common.data.Device;
 import org.thingsboard.server.common.data.DeviceProfileProvisionType;
 import org.thingsboard.server.common.data.TransportPayloadType;
 import org.thingsboard.server.common.data.device.credentials.BasicMqttCredentials;
-import org.thingsboard.server.common.data.device.profile.MqttTopics;
 import org.thingsboard.server.common.data.security.DeviceCredentials;
 import org.thingsboard.server.common.data.security.DeviceCredentialsType;
 import org.thingsboard.server.common.msg.EncryptionUtil;
@@ -46,11 +41,15 @@ import org.thingsboard.server.gen.transport.TransportProtos.ValidateBasicMqttCre
 import org.thingsboard.server.gen.transport.TransportProtos.ValidateDeviceTokenRequestMsg;
 import org.thingsboard.server.gen.transport.TransportProtos.ValidateDeviceX509CertRequestMsg;
 import org.thingsboard.server.transport.mqtt.AbstractMqttIntegrationTest;
+import org.thingsboard.server.transport.mqtt.MqttTestCallback;
+import org.thingsboard.server.transport.mqtt.MqttTestClient;
 import org.thingsboard.server.transport.mqtt.MqttTestConfigProperties;
 
-import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.TimeUnit;
 
+import static org.thingsboard.server.common.data.device.profile.MqttTopics.DEVICE_PROVISION_REQUEST_TOPIC;
+import static org.thingsboard.server.common.data.device.profile.MqttTopics.DEVICE_PROVISION_RESPONSE_TOPIC;
+
 @Slf4j
 @DaoSqlTest
 public class MqttProvisionProtoDeviceTest extends AbstractMqttIntegrationTest {
@@ -104,9 +103,9 @@ public class MqttProvisionProtoDeviceTest extends AbstractMqttIntegrationTest {
                 .provisionType(DeviceProfileProvisionType.DISABLED)
                 .build();
         processBeforeTest(configProperties);
-        ProvisionDeviceResponseMsg result = ProvisionDeviceResponseMsg.parseFrom(createMqttClientAndPublish().getPayloadBytes());
+        ProvisionDeviceResponseMsg result = ProvisionDeviceResponseMsg.parseFrom(createMqttClientAndPublish());
         Assert.assertNotNull(result);
-        Assert.assertEquals(ProvisionResponseStatus.NOT_FOUND.name(), result.getStatus().toString());
+        Assert.assertEquals(ProvisionResponseStatus.NOT_FOUND.name(), result.getStatus().name());
     }
 
     protected void processTestProvisioningCreateNewDeviceWithoutCredentials() throws Exception {
@@ -118,7 +117,7 @@ public class MqttProvisionProtoDeviceTest extends AbstractMqttIntegrationTest {
                 .provisionSecret("testProvisionSecret")
                 .build();
         processBeforeTest(configProperties);
-        ProvisionDeviceResponseMsg response = ProvisionDeviceResponseMsg.parseFrom(createMqttClientAndPublish().getPayloadBytes());
+        ProvisionDeviceResponseMsg response = ProvisionDeviceResponseMsg.parseFrom(createMqttClientAndPublish());
 
         Device createdDevice = deviceService.findDeviceByTenantIdAndName(tenantId, "Test Provision device");
 
@@ -126,8 +125,8 @@ public class MqttProvisionProtoDeviceTest extends AbstractMqttIntegrationTest {
 
         DeviceCredentials deviceCredentials = deviceCredentialsService.findDeviceCredentialsByDeviceId(tenantId, createdDevice.getId());
 
-        Assert.assertEquals(deviceCredentials.getCredentialsType().name(), response.getCredentialsType().toString());
+        Assert.assertEquals(deviceCredentials.getCredentialsType().name(), response.getCredentialsType().name());
-        Assert.assertEquals(ProvisionResponseStatus.SUCCESS.name(), response.getStatus().toString());
+        Assert.assertEquals(ProvisionResponseStatus.SUCCESS.name(), response.getStatus().name());
     }
 
     protected void processTestProvisioningCreateNewDeviceWithAccessToken() throws Exception {
@@ -139,9 +138,12 @@ public class MqttProvisionProtoDeviceTest extends AbstractMqttIntegrationTest {
                 .provisionSecret("testProvisionSecret")
                 .build();
         processBeforeTest(configProperties);
-        CredentialsDataProto requestCredentials = CredentialsDataProto.newBuilder().setValidateDeviceTokenRequestMsg(ValidateDeviceTokenRequestMsg.newBuilder().setToken("test_token").build()).build();
+        CredentialsDataProto requestCredentials = CredentialsDataProto.newBuilder()
+                .setValidateDeviceTokenRequestMsg(ValidateDeviceTokenRequestMsg.newBuilder().setToken("test_token").build())
+                .build();
 
-        ProvisionDeviceResponseMsg response = ProvisionDeviceResponseMsg.parseFrom(createMqttClientAndPublish(createTestsProvisionMessage(CredentialsType.ACCESS_TOKEN, requestCredentials)).getPayloadBytes());
+        ProvisionDeviceResponseMsg response = ProvisionDeviceResponseMsg.parseFrom(
+                createMqttClientAndPublish(createTestsProvisionMessage(CredentialsType.ACCESS_TOKEN, requestCredentials)));
 
         Device createdDevice = deviceService.findDeviceByTenantIdAndName(tenantId, "Test Provision device");
 
@@ -164,9 +166,13 @@ public class MqttProvisionProtoDeviceTest extends AbstractMqttIntegrationTest {
                 .provisionSecret("testProvisionSecret")
                 .build();
         processBeforeTest(configProperties);
-        CredentialsDataProto requestCredentials = CredentialsDataProto.newBuilder().setValidateDeviceX509CertRequestMsg(ValidateDeviceX509CertRequestMsg.newBuilder().setHash("testHash").build()).build();
+        CredentialsDataProto requestCredentials = CredentialsDataProto.newBuilder()
+                .setValidateDeviceX509CertRequestMsg(
+                        ValidateDeviceX509CertRequestMsg.newBuilder().setHash("testHash").build())
+                .build();
 
-        ProvisionDeviceResponseMsg response = ProvisionDeviceResponseMsg.parseFrom(createMqttClientAndPublish(createTestsProvisionMessage(CredentialsType.X509_CERTIFICATE, requestCredentials)).getPayloadBytes());
+        ProvisionDeviceResponseMsg response = ProvisionDeviceResponseMsg.parseFrom(
+                createMqttClientAndPublish(createTestsProvisionMessage(CredentialsType.X509_CERTIFICATE, requestCredentials)));
 
         Device createdDevice = deviceService.findDeviceByTenantIdAndName(tenantId, "Test Provision device");
 
@@ -203,7 +209,8 @@ public class MqttProvisionProtoDeviceTest extends AbstractMqttIntegrationTest {
                     .build()
         ).build();
 
-        ProvisionDeviceResponseMsg response = ProvisionDeviceResponseMsg.parseFrom(createMqttClientAndPublish(createTestsProvisionMessage(CredentialsType.MQTT_BASIC, requestCredentials)).getPayloadBytes());
+        ProvisionDeviceResponseMsg response = ProvisionDeviceResponseMsg.parseFrom(
+                createMqttClientAndPublish(createTestsProvisionMessage(CredentialsType.MQTT_BASIC, requestCredentials)));
 
         Device createdDevice = deviceService.findDeviceByTenantIdAndName(tenantId, "Test Provision device");
 
@@ -233,12 +240,12 @@ public class MqttProvisionProtoDeviceTest extends AbstractMqttIntegrationTest {
                 .provisionSecret("testProvisionSecret")
                 .build();
         processBeforeTest(configProperties);
-        ProvisionDeviceResponseMsg response = ProvisionDeviceResponseMsg.parseFrom(createMqttClientAndPublish().getPayloadBytes());
+        ProvisionDeviceResponseMsg response = ProvisionDeviceResponseMsg.parseFrom(createMqttClientAndPublish());
 
         DeviceCredentials deviceCredentials = deviceCredentialsService.findDeviceCredentialsByDeviceId(tenantId, savedDevice.getId());
 
-        Assert.assertEquals(deviceCredentials.getCredentialsType().name(), response.getCredentialsType().toString());
+        Assert.assertEquals(deviceCredentials.getCredentialsType().name(), response.getCredentialsType().name());
-        Assert.assertEquals(ProvisionResponseStatus.SUCCESS.name(), response.getStatus().toString());
+        Assert.assertEquals(ProvisionResponseStatus.SUCCESS.name(), response.getStatus().name());
     }
 
     protected void processTestProvisioningWithBadKeyDevice() throws Exception {
@@ -250,70 +257,25 @@ public class MqttProvisionProtoDeviceTest extends AbstractMqttIntegrationTest {
                 .provisionSecret("testProvisionSecret")
                 .build();
         processBeforeTest(configProperties);
-        ProvisionDeviceResponseMsg response = ProvisionDeviceResponseMsg.parseFrom(createMqttClientAndPublish().getPayloadBytes());
+        ProvisionDeviceResponseMsg response = ProvisionDeviceResponseMsg.parseFrom(createMqttClientAndPublish());
-        Assert.assertEquals(ProvisionResponseStatus.NOT_FOUND.name(), response.getStatus().toString());
+        Assert.assertEquals(ProvisionResponseStatus.NOT_FOUND.name(), response.getStatus().name());
     }
 
-    protected TestMqttCallback createMqttClientAndPublish() throws Exception {
+    protected byte[] createMqttClientAndPublish() throws Exception {
         byte[] provisionRequestMsg = createTestProvisionMessage();
         return createMqttClientAndPublish(provisionRequestMsg);
     }
 
-    protected TestMqttCallback createMqttClientAndPublish(byte[] provisionRequestMsg) throws Exception {
+    protected byte[] createMqttClientAndPublish(byte[] provisionRequestMsg) throws Exception {
-        MqttAsyncClient client = getMqttAsyncClient("provision");
+        MqttTestClient client = new MqttTestClient();
-        TestMqttCallback onProvisionCallback = getTestMqttCallback();
+        client.connectAndWait("provision");
+        MqttTestCallback onProvisionCallback = new MqttTestCallback(DEVICE_PROVISION_RESPONSE_TOPIC);
         client.setCallback(onProvisionCallback);
-        client.subscribe(MqttTopics.DEVICE_PROVISION_RESPONSE_TOPIC, MqttQoS.AT_MOST_ONCE.value());
+        client.subscribe(DEVICE_PROVISION_RESPONSE_TOPIC, MqttQoS.AT_MOST_ONCE);
-        Thread.sleep(2000);
+        client.publishAndWait(DEVICE_PROVISION_REQUEST_TOPIC, provisionRequestMsg);
-        client.publish(MqttTopics.DEVICE_PROVISION_REQUEST_TOPIC, new MqttMessage(provisionRequestMsg));
+        onProvisionCallback.getSubscribeLatch().await(3, TimeUnit.SECONDS);
-        onProvisionCallback.getLatch().await(3, TimeUnit.SECONDS);
+        client.disconnect();
-        return onProvisionCallback;
+        return onProvisionCallback.getPayloadBytes();
-    }
-
-
-    protected TestMqttCallback getTestMqttCallback() {
-        CountDownLatch latch = new CountDownLatch(1);
-        return new TestMqttCallback(latch);
-    }
-
-
-    protected static class TestMqttCallback implements MqttCallback {
-
-        private final CountDownLatch latch;
-        private Integer qoS;
-        private byte[] payloadBytes;
-
-        TestMqttCallback(CountDownLatch latch) {
-            this.latch = latch;
-        }
-
-        public int getQoS() {
-            return qoS;
-        }
-
-        public byte[] getPayloadBytes() {
-            return payloadBytes;
-        }
-
-        public CountDownLatch getLatch() {
-            return latch;
-        }
-
-        @Override
-        public void connectionLost(Throwable throwable) {
-        }
-
-        @Override
-        public void messageArrived(String requestTopic, MqttMessage mqttMessage) throws Exception {
-            qoS = mqttMessage.getQos();
-            payloadBytes = mqttMessage.getPayload();
-            latch.countDown();
-        }
-
-        @Override
-        public void deliveryComplete(IMqttDeliveryToken iMqttDeliveryToken) {
-
-        }
     }
 
     protected byte[] createTestsProvisionMessage(CredentialsType credentialsType, CredentialsDataProto credentialsData) throws Exception {