def test_not_revoke_cert(self, mock_utils,
mock_logging,
mock_cert_client,
paho_mqtt_mock,
json_mock):
"Sould not revoke cert"
client = MQTTClient("123", "987", True, False)
client.should_revoke_now = MagicMock()
client.should_revoke_now.return_value = True
mock_cert_client.revoke_cert = None
client.setup()
with self.assertRaises(Exception):
client.revoke_cert()
mock_utils.fire_locust_failure.assert_called_once()
client.should_revoke_now.return_value = False
client.revoke_cert()
mock_utils.reset_mock()
mock_logging.reset_mock()
mock_cert_client.reset_mock()
paho_mqtt_mock.reset_mock()
json_mock.reset_mock()
class MqttLocust(Locust):
"""Locust client using MQTT."""
def __init__(self, *args, **kwargs):
super(Locust, self).__init__(*args, **kwargs)
# Connects to Redis database that stores the device_id for each client
cache = RedisClient()
revoke = cache.has_to_revoke()
should_revoke = False
should_renew = False
device_id = None
# We need to differentiate the device IDs to be revogated/renewed from the other ones
# The revogated/renewed ones will not be stored in Redis; instead, they will be created
# at runtime
if revoke:
should_revoke = revoke["should_revoke"]
device_id = revoke["device_id"]
else:
renew = cache.has_to_renew()
if renew:
should_renew = renew["should_renew"]
device_id = renew["device_id"]
else:
device_id = cache.next_device_id()
# UUID to identify the client run
run_id = str(uuid.uuid4())
self.client = MQTTClient(device_id, run_id, should_revoke,
should_renew)
self.client.setup()
self.client.connect()
def test_revoke_cert(self, mock_utils,
mock_logging,
mock_cert_client,
paho_mqtt_mock,
json_mock):
"Sould not revoke cert"
client = MQTTClient("123", "987", True, False)
client.should_revoke_now = MagicMock()
client.should_revoke_now.return_value = True
mock_cert_client.revoke_cert.return_value = MagicMock()
mock_cert_client.has_been_revoked.return_value = True
client.setup()
client.revoke_cert()
mock_cert_client.revoke_cert.assert_called_once()
mock_cert_client.has_been_revoked.assert_called_once()
mock_utils.fire_locust_success.assert_called_once()
mock_cert_client.has_been_revoked.return_value = False
client.revoke_cert()
mock_utils.fire_locust_success.assert_called_once()
mock_utils.reset_mock()
mock_logging.reset_mock()
mock_cert_client.reset_mock()
paho_mqtt_mock.reset_mock()
json_mock.reset_mock()
def test_should_revoke_now(self, mock_utils,
mock_logging,
mock_cert_client,
paho_mqtt_mock,
json_mock):
"""
Should revoke the certificate now.
"""
mock_utils.time_delta.return_value = 1202010
client = MQTTClient("123", "987", True, False)
client.setup()
self.assertTrue(client.should_revoke_now())
mock_utils.reset_mock()
mock_logging.reset_mock()
mock_cert_client.reset_mock()
paho_mqtt_mock.reset_mock()
json_mock.reset_mock()
def test_setup_success(self, mock_utils,
mock_logging,
mock_cert_client,
paho_mqtt_mock,
json_mock):
"""
Should setup the parameters successfully.
"""
client = MQTTClient("123", "987", False, False)
client.configure_mqtt = MagicMock()
client.setup()
mock_logging.basicConfig.assert_called_once()
client.configure_mqtt.assert_called_once()
mock_utils.reset_mock()
mock_logging.reset_mock()
mock_cert_client.reset_mock()
paho_mqtt_mock.reset_mock()
json_mock.reset_mock()
def test_locust_on_disconnect_fail(self, mock_utils,
mock_logging,
mock_cert_client,
paho_mqtt_mock,
json_mock):
"""
Should fire locust failure on disconnect callback
"""
paho_mqtt_mock.MQTT_ERR_SUCCESS = 1
client = MQTTClient("123", "987", False, False)
client.setup()
client.locust_on_disconnect(client.mqttc, {}, 1)
self.assertFalse(client.is_connected)
paho_mqtt_mock.Client().reconnect.assert_called_once()
mock_utils.reset_mock()
mock_logging.reset_mock()
mock_cert_client.reset_mock()
paho_mqtt_mock.reset_mock()
json_mock.reset_mock()
def test_subcribe_succes(self, mock_utils,
mock_logging,
mock_cert_client,
paho_mqtt_mock,
json_mock):
"""
Should subscribe a message successfully
"""
paho_mqtt_mock.Client().subscribe.return_value = (0, MagicMock())
client = MQTTClient("123", "987", False, False)
client.setup()
client.subscribing()
paho_mqtt_mock.Client().subscribe.assert_called_once()
keys_len = len(client.submmap.keys())
self.assertGreater(keys_len, 0)
mock_utils.reset_mock()
mock_logging.reset_mock()
mock_cert_client.reset_mock()
paho_mqtt_mock.reset_mock()
json_mock.reset_mock()
def test_connect_error(self, mock_utils,
mock_logging,
mock_cert_client,
paho_mqtt_mock,
json_mock):
"""
Should not connect to the broker successfully - connect_async error.
"""
paho_mqtt_mock.Client().connect_async = Exception("fake error")
client = MQTTClient("123", "987", False, False)
client.setup()
# when exception raises
client.connect()
mock_utils.fire_locust_failure.assert_called_once()
mock_utils.reset_mock()
mock_logging.reset_mock()
mock_cert_client.reset_mock()
paho_mqtt_mock.reset_mock()
json_mock.reset_mock()
def test_setup_success_renew(self, mock_utils,
mock_logging,
mock_cert_client,
paho_mqtt_mock,
json_mock):
"""
Should setup the parameters successfully with certificate renovation.
"""
client = MQTTClient("123", "987", False, True)
client.configure_mqtt = MagicMock()
client.setup()
mock_cert_client.new_cert.assert_called_once_with(
client.tenant, client.device_id)
mock_cert_client.create_cert_files.assert_called_once_with(
client.new_cert, client.device_cert_dir)
mock_utils.reset_mock()
mock_logging.reset_mock()
mock_cert_client.reset_mock()
paho_mqtt_mock.reset_mock()
json_mock.reset_mock()
def test_renew_cert(self, mock_utils,
mock_logging,
mock_cert_client,
paho_mqtt_mock,
json_mock):
"Sould renew cert"
client = MQTTClient("123", "987", False, True)
client.should_renew_now = MagicMock()
client.should_renew_now.return_value = True
mock_cert_client.new_cert().renew_cert.return_value = True
client.setup()
client.renew_cert()
client.new_cert.renew_cert.assert_called_once()
mock_utils.fire_locust_success.assert_called_once()
self.assertFalse(client.should_renew)
mock_utils.reset_mock()
mock_logging.reset_mock()
mock_cert_client.reset_mock()
paho_mqtt_mock.reset_mock()
json_mock.reset_mock()
def test_subscribe_error(self, mock_utils,
mock_logging,
mock_cert_client,
paho_mqtt_mock,
json_mock):
"""
Should not subscribe a message successfully
"""
paho_mqtt_mock.Client().subscribe.return_value = (10, MagicMock())
client = MQTTClient("123", "987", False, False)
client.setup()
client.subscribing()
mock_utils.error_message.assert_called_once_with(10)
mock_utils.fire_locust_failure.assert_called_once()
keys_len = len(client.pubmmap.keys())
self.assertEqual(keys_len, 0)
mock_utils.reset_mock()
mock_logging.reset_mock()
mock_cert_client.reset_mock()
paho_mqtt_mock.reset_mock()
json_mock.reset_mock()
def test_connect(self, mock_utils,
mock_logging,
mock_cert_client,
paho_mqtt_mock,
json_mock):
"""
Should connect to the broker successfully.
"""
client = MQTTClient("123", "987", False, False)
client.setup()
client.connect()
paho_mqtt_mock.Client().connect_async.assert_called_once_with(
host=MOCK_CONFIG['mqtt']['host'],
port=MOCK_CONFIG['mqtt']['port'],
keepalive=MOCK_CONFIG['mqtt']['con_timeout'])
paho_mqtt_mock.Client().loop_start.assert_called_once()
mock_utils.reset_mock()
mock_logging.reset_mock()
mock_cert_client.reset_mock()
paho_mqtt_mock.reset_mock()
json_mock.reset_mock()
