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_connect_error(self, _mock_paho, mock_utils):
"""
Should not connect to the broker successfully - connect_async error.
"""
client = MQTTClient("123", "987", False, False)
client.mqttc.connect_async.side_effect = Exception
client.connect()
self.assertRaises(Exception, client.mqttc.connect_async)
client.mqttc.loop_start.assert_not_called()
mock_utils.fire_locust_failure.assert_called_once()
def test_connect(self, mock_paho, _mock_utils):
"""
Should connect to the broker successfully.
"""
client = MQTTClient("123", "987", False, False)
client.connect()
mock_paho.Client().connect_async.assert_called_once_with(
host=MOCK_CONFIG['mqtt']['host'],
port=MOCK_CONFIG['mqtt']['port'],
keepalive=MOCK_CONFIG['mqtt']['con_timeout'])
mock_paho.Client().loop_start.assert_called_once()
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()
