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