def test_create_pubsub_message( force, ping, top_down, version, expected_klio_job ): entity_id = "s0m3-ent1ty-1D" expected_klio_message = klio_pb2.KlioMessage() expected_klio_message.metadata.force = force expected_klio_message.metadata.ping = ping expected_klio_message.version = version if version == 1: expected_klio_message.data.entity_id = entity_id if not top_down: expected_klio_message.metadata.downstream.extend( [expected_klio_job] ) else: expected_klio_message.data.element = bytes(entity_id, "utf-8") if not top_down: rec = expected_klio_message.metadata.intended_recipients rec.limited.recipients.extend([expected_klio_job]) ret_msg = publish._create_pubsub_message( entity_id, expected_klio_job, force, ping, top_down, version ) assert expected_klio_message.SerializeToString() == ret_msg
def _get_klio_message(): parent_klio_job = _get_klio_job() msg = klio_pb2.KlioMessage() msg.metadata.visited.extend([parent_klio_job]) msg.metadata.force = True msg.metadata.ping = True msg.data.element = b"1234567890" msg.version = klio_pb2.Version.V2 return msg
def test_from_klio_message_v1(): payload = b"some-payload" msg = klio_pb2.KlioMessage() msg.version = klio_pb2.Version.V1 msg.data.payload = payload expected_str = msg.SerializeToString() actual_message = serializer.from_klio_message(msg, payload) assert expected_str == actual_message
def test_to_klio_message_allow_non_kmsg(klio_config, logger, monkeypatch): monkeypatch.setattr( klio_config.job_config, "allow_non_klio_messages", True ) incoming = b"Not a klio message" expected = klio_pb2.KlioMessage() expected.data.element = incoming expected.version = klio_pb2.Version.V2 actual_message = serializer.to_klio_message(incoming, klio_config, logger) assert expected == actual_message logger.error.assert_not_called()
def test_handle_msg_compat(version, element, entity_id, payload): msg = klio_pb2.KlioMessage() msg.version = version if element: msg.data.element = element if payload: msg.data.payload = payload if entity_id: msg.data.entity_id = entity_id actual_msg = serializer._handle_msg_compat(msg) assert actual_msg.version is not klio_pb2.Version.UNKNOWN # we assume in the function's logic that v2 messages are already parsed # correctly if entity_id and not klio_pb2.Version.V2: assert entity_id == actual_msg.data.element.decode("utf-8")
def _create_pubsub_message(entity_id, job, force, ping, top_down, msg_version): kmsg = klio_pb2.KlioMessage() kmsg.version = msg_version if msg_version == 1: kmsg.data.entity_id = entity_id if not top_down: kmsg.metadata.downstream.extend([job]) elif msg_version == 2: kmsg.data.element = bytes(entity_id, "utf-8") if not top_down: kmsg.metadata.intended_recipients.limited.recipients.extend([job]) else: kmsg.metadata.intended_recipients.anyone.SetInParent() kmsg.metadata.ping = ping kmsg.metadata.force = force return kmsg.SerializeToString()
def test_private_publish_messages( klio_job_config, mock_publisher, expected_klio_job, force, ping, top_down, non_klio, ): entity_id = "s0m3-ent1ty-1D" msg_version = 1 if non_klio: exp_data = bytes(entity_id.encode("utf-8")) else: expected_klio_message = klio_pb2.KlioMessage() expected_klio_message.metadata.force = force expected_klio_message.metadata.ping = ping expected_klio_message.data.entity_id = entity_id expected_klio_message.version = msg_version if not top_down: expected_klio_message.metadata.downstream.extend( [expected_klio_job] ) exp_data = expected_klio_message.SerializeToString() ret_success, ret_fail = publish._publish_messages( klio_job_config, [entity_id], force, ping, top_down, non_klio, msg_version, ) mock_publisher.return_value.publish.assert_called_once_with( topic="an-input-topic", data=exp_data ) assert 1 == len(ret_success) assert not len(ret_fail)
def expected_klio_message(expected_klio_job): message = klio_pb2.KlioMessage() message.metadata.downstream.extend([expected_klio_job]) return message