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_locust_on_publish_failure(self, mock_utils, _mock_paho): """ Should fire locust failure on publish """ client = MQTTClient("123", "987", False, False) client.locust_on_publish(client.mqttc, {}, 0) mock_utils.fire_locust_failure.assert_called_once()
def test_configure_mqtt(self, mock_cert_utils, mock_paho): """ Should configure the MQTT connection. """ client = MQTTClient("123", "987", False, True) # configure_mqtt() function is called at the constructor, so # we need to reset the mocks to be able to test it mock_cert_utils.reset_mock() mock_paho.reset_mock() client.configure_mqtt() mock_cert_utils.get_certificate_file.assert_called_once_with( client.device_id) mock_cert_utils.get_private_key_file.assert_called_once_with( client.device_id) self.assertIsNotNone(client.mqttc) mock_paho.Client.assert_called_once_with(client_id=client.device_id) client.mqttc.reconnect_delay_set.assert_called_once_with( min_delay=MOCK_CONFIG["security"]["min_time_reconn"], max_delay=MOCK_CONFIG["security"]["max_time_reconn"]) client.mqttc.tls_set.assert_called_once() client.mqttc.tls_insecure_set.assert_called_once_with(True) self.assertIsNotNone(client.mqttc.on_connect) self.assertIsNotNone(client.mqttc.on_disconnect) self.assertIsNotNone(client.mqttc.on_publish) self.assertIsNotNone(client.mqttc.on_subscribe) self.assertIsNotNone(client.mqttc.on_message)
def test_locust_on_subcribing_failure(self, mock_utils, _mock_paho): """ Should fire locust failure """ client = MQTTClient("123", "987", False, False) client.locust_on_subscribe(client.mqttc, {}, 0, 0) mock_utils.fire_locust_failure.assert_called_once()
def test_should_not_renew_now(self, _mock_utils, _mock_paho, _mock_cert_utils): """ Should not renew the certificate now. """ client = MQTTClient("123", "987", False, False) self.assertFalse(client.should_renew_now())
def test_should_revoke_now(self, mock_utils, _mock_paho, _mock_cert_utils): """ Should revoke the certificate now. """ client = MQTTClient("123", "987", True, False) mock_utils.seconds_to_milliseconds.return_value = 200 self.assertTrue(client.should_revoke_now())
def test_locust_on_subcribing(self, mock_utils, _mock_paho): """ Should fire locust success """ mid = MagicMock() client = MQTTClient("123", "987", False, False) client.submmap[mid] = {'name': 'name', 'start_time': 'time'} client.locust_on_subscribe(client.mqttc, {}, mid, 0) mock_utils.fire_locust_success.assert_called_once()
def test_disconnect(self, _mock_paho): """ Should disconnect correctly. """ client = MQTTClient("123", "987", False, False) client.disconnect() client.mqttc.disconnect.assert_called_once()
def test_locust_on_connect_failure(self, mock_utils, mock_paho): """ Should fire locust failure on connection callback """ mock_paho.MQTT_ERR_SUCCESS = 1 client = MQTTClient("123", "987", False, False) client.locust_on_connect(client.mqttc, {}, {}, 101010) mock_utils.error_message.assert_called_once() mock_utils.fire_locust_failure.assert_called_once()
def test_locust_on_connect(self, mock_utils, mock_paho): """ Should fire locust success on connection callback """ mock_paho.MQTT_ERR_SUCCESS = 1 client = MQTTClient("123", "987", False, False) client.subscribe = MagicMock() client.locust_on_connect(client.mqttc, {}, {}, 1) client.subscribe.assert_called_once() mock_utils.fire_locust_success.assert_called_once()
def test_locust_on_disconnect_fail(self, _mock_utils, mock_paho): """ Should fire locust failure on disconnect callback """ mock_paho.MQTT_ERR_SUCCESS = 1 client = MQTTClient("123", "987", False, False) client.locust_on_disconnect(client.mqttc, {}, 1) self.assertFalse(client.is_connected) mock_paho.Client().reconnect.assert_called_once()
def test_reconnect(self, _mock_json, _mock_paho, _mock_utils, _mock_cert_utils): """ Should reconnect successfully. """ client = MQTTClient("123", "987", False, False) client.reconnect() client.mqttc.reconnect.assert_called_once()
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()
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_subcribe_success(self, _mock_utils, mock_paho): """ Should subscribe to the topic successfully. """ mock_paho.Client().subscribe.return_value = (0, MagicMock()) client = MQTTClient("123", "987", False, False) client.subscribe() mock_paho.Client().subscribe.assert_called_once() keys_len = len(client.submmap.keys()) self.assertGreater(keys_len, 0)
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_publish_success(self, _mock_utils, mock_paho, _mock_json): """ Should publish a message successfully """ mock_paho.Client().publish.return_value = (0, MagicMock()) client = MQTTClient("123", "987", False, False) client.publish() mock_paho.Client().publish.assert_called_once() keys_len = len(client.pubmmap.keys()) self.assertGreater(keys_len, 0)
def test_subscribe_error(self, mock_utils, mock_paho): """ Should not subscribe a message successfully """ mock_paho.Client().subscribe.return_value = (10, MagicMock()) client = MQTTClient("123", "987", False, False) client.subscribe() 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)
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_locust_on_message(self, mock_utils, _mock_paho, mock_json): """ Should fire locust sucess on message callback """ message: mqtt.MQTTMessage = mqtt.MQTTMessage() message.payload = str.encode(str({"timestamp": 0})) mock_json.loads.return_value = {"timestamp": 0} client = MQTTClient("123", "987", False, False) client.locust_on_message(client.mqttc, {}, message) mock_utils.fire_locust_success.assert_called_once()
def test_locust_on_publish(self, mock_utils, _mock_paho): """ Should fire locust success on publish callback """ mid = MagicMock() client = MQTTClient("123", "987", False, False) client.pubmmap[mid] = { 'name': 'name', 'start_time': 'time', 'payload': 'payload' } client.locust_on_publish(client.mqttc, {}, mid) mock_utils.fire_locust_success.assert_called_once()
def test_revoke_cert_and_emit_event(self, mock_utils, _mock_paho, mock_cert_utils): """ Should revoke cert """ client = MQTTClient("123", "987", True, False) mock_cert_utils.has_been_revoked.return_value = False has_revoked = client.revoke_cert_and_emit_event() self.assertTrue(has_revoked) mock_cert_utils.revoke_cert.assert_called_once() mock_utils.fire_locust_success.assert_called()
def test_not_revoke_cert_and_emit_event(self, mock_utils, _mock_paho, mock_cert_utils): """ Should not revoke cert - Exception thrown by the revoking function """ client = MQTTClient("123", "987", True, False) mock_cert_utils.has_been_revoked.return_value = False mock_cert_utils.revoke_cert.side_effect = Exception has_revoked = client.revoke_cert_and_emit_event() self.assertFalse(has_revoked) mock_cert_utils.revoke_cert.assert_called_once() mock_utils.fire_locust_failure.assert_called()
def test_renew_cert_and_emit_event(self, mock_utils, _mock_paho, mock_cert_utils): "Should renew cert" client = MQTTClient("123", "987", False, True) client.revoke_cert_and_emit_event = MagicMock() client.revoke_cert_and_emit_event.return_value = True mock_cert_utils.create_cert_files = MagicMock() has_renewed = client.renew_cert_and_emit_event() client.revoke_cert_and_emit_event.assert_called_once() client.new_cert.renew_cert.assert_called_once() mock_cert_utils.create_cert_files.assert_called_once() mock_utils.fire_locust_success.assert_called_once() self.assertTrue(has_renewed)
def test_constructor_both_true(self, _mock_utils, _mock_paho): """ Should raise a ValueError when passing True to both should_revoke and should_renew. """ with self.assertRaises(ValueError): MQTTClient("123", "987", True, True)
def test_not_renew_cert_and_emit_event(self, mock_utils, _mock_paho, mock_cert_utils): "Should not renew cert" client = MQTTClient("123", "987", False, True) client.revoke_cert_and_emit_event = MagicMock() client.revoke_cert_and_emit_event.return_value = False mock_cert_utils.create_cert_files = MagicMock() has_renewed = client.renew_cert_and_emit_event() client.revoke_cert_and_emit_event.assert_called_once() self.assertTrue(client.new_cert.renew_cert.call_count == 0) self.assertTrue(mock_cert_utils.create_cert_files.call_count == 0) mock_utils.fire_locust_failure.assert_called() self.assertFalse(has_renewed)
def test_constructor_success(self, mock_utils, _mock_paho): """ Should create a MQTTClient instance. """ device_id = "123" run_id = "987" tenant = MOCK_CONFIG['app']['tenant'] client = MQTTClient(device_id, run_id, False, False) mock_utils.validate_tenant.assert_called_once_with(tenant) mock_utils.validate_device_id.assert_called_once_with(device_id) self.assertEqual(client.device_id, device_id) self.assertEqual(client.run_id, run_id) self.assertEqual(client.should_revoke, False) self.assertEqual(client.should_renew, False) self.assertFalse(client.is_connected) self.assertEqual(client.disconnect_forever, False) self.assertIsNotNone(client.mqttc) self.assertEqual(client.tenant, tenant) self.assertIsNotNone(client.username) self.assertIsNotNone(client.topic) self.assertIsNotNone(client.sub_topic) self.assertEqual(client.device_cert_dir, MOCK_CONFIG["security"]["cert_dir"]) self.assertIsNone(client.new_cert) self.assertEqual(client.pubmmap, {}) self.assertEqual(client.submmap, {}) self.assertEqual(client.start_time, 0)
def test_locust_on_publish_failure(self, mock_utils, mock_logging, mock_cert_client, paho_mqtt_mock, json_mock): """ Should fire locust failure on publish """ client = MQTTClient("123", "987", False, False) client.locust_on_publish(client.mqttc, {}, 0) 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_renew_cert_and_emit_event_exception(self, mock_utils, _mock_paho, mock_cert_utils): "Should raise an exception when renewing cert" client = MQTTClient("123", "987", False, True) client.revoke_cert_and_emit_event = MagicMock() client.revoke_cert_and_emit_event.return_value = True mock_cert_utils.create_cert_files = MagicMock() client.new_cert.renew_cert.side_effect = Exception has_renewed = client.renew_cert_and_emit_event() client.revoke_cert_and_emit_event.assert_called_once() self.assertRaises(Exception, client.new_cert.renew_cert) self.assertTrue(mock_cert_utils.create_cert_files.call_count == 0) mock_utils.fire_locust_failure.assert_called() self.assertFalse(has_renewed)
def test_create_certificate_renew(self, mock_cert_utils, _mock_paho): """ create_certificate() should create the certificates for renewal """ MQTTClient("123", "987", True, False) mock_cert_utils.new_cert.assert_called_once() mock_cert_utils.create_cert_files.assert_called_once()