def _try_extract_egress(self, invocation_result):
if len(invocation_result.outgoing_egresses) == 0:
return None
elif len(invocation_result.outgoing_egresses) > 1:
raise ValueError(
f'Expected at most 1 egress message. Found {len(invocation_result.outgoing_egresses)}'
)
kafka_producer_record = KafkaProducerRecord()
kafka_producer_record.ParseFromString(
invocation_result.outgoing_egresses[0].argument.value)
result_any = Any.FromString(kafka_producer_record.value_bytes)
result_proto = unpack_any(
result_any,
[TaskResult, TaskException, TaskActionResult, TaskActionException])
if isinstance(result_proto, TaskResult):
result, _ = test_harness_serialiser.deserialise_result(
result_proto)
return result
elif isinstance(result_proto, TaskException):
return TaskErrorException(TaskError(result_proto))
elif isinstance(result_proto, TaskActionException):
return TaskErrorException(TaskError(result_proto))
elif isinstance(result_proto, TaskActionResult):
return result_proto
def handle_log_visualize(self):
"""We send protobufs from our C++ code to python for visualization.
If we used the handle_proto handler and passed in a proto_class, we
would need to setup a sender/receiver pair for every protobuf we want
to visualize.
So instead, we special case the communication coming from the ProtobufSink
(C++ side). The ProtobufSink sends the typename prefixed at the beginning
of the payload delimited by the TYPE_DELIMITER (!!!).
| -- data -- |
PackageName.TypeName!!!eW91Zm91bmR0aGVzZWNyZXRtZXNzYWdl
This allows us to call LOG(VISUALIZE) _anywhere_ in C++ and
receive/decode here with minimum boilerplate code.
"""
payload = self.request[0]
type_name = str(payload.split(b"!!!")[0], "utf-8")
proto_type = self.find_proto_class(type_name.split(".")[1])
msg = proto_type()
payload = base64.b64decode(payload.split(b"!!!")[1])
any_msg = Any.FromString(payload)
any_msg.Unpack(msg)
self.handle_callback(msg)
def _try_extract_egress(self, invocation_result):
if len(invocation_result.outgoing_egresses) == 0:
return None
elif len(invocation_result.outgoing_egresses) > 1:
raise ValueError(f'Expected at most 1 egress message. Found {len(invocation_result.outgoing_egresses)}')
kafka_producer_record = KafkaProducerRecord()
invocation_result.outgoing_egresses[0].argument.Unpack(kafka_producer_record)
if kafka_producer_record.topic != self.__reply_topic:
raise ValueError(f'Unexpected topic for egress: {kafka_producer_record.topic}')
result_any = Any.FromString(kafka_producer_record.value_bytes)
task_result_or_exception = unpack_any(result_any, [TaskResult, TaskException])
if isinstance(task_result_or_exception, TaskResult):
result, _ = serialiser.deserialise_result(task_result_or_exception)
return result
else:
raise TaskErrorException(TaskError(task_result_or_exception))
def handle(self):
"""Decode the base64 request and unpack from Any if we are receiving
an Any protobuf. If not, just unpack directly into the type provided.
Then, trigger the handle callback
"""
p = base64.b64decode(self.request[0])
msg = self.proto_type()
if self.convert_from_any:
any_msg = Any.FromString(p)
any_msg.Unpack(msg)
else:
msg = self.proto_type.FromString(p)
self.handle_callback(msg)
