def test_proto_to_literal_type():
proto_type = _proto.create_protobuf(_errors_pb2.ContainerError)
assert proto_type.to_flyte_literal_type(
).simple == _type_models.SimpleType.BINARY
assert len(proto_type.to_flyte_literal_type().metadata) == 1
assert proto_type.to_flyte_literal_type().metadata[_proto.Protobuf.PB_FIELD_KEY] == \
"flyteidl.core.errors_pb2.ContainerError"
def test_from_string():
proto_type = _proto.create_protobuf(_errors_pb2.ContainerError)
pb = _errors_pb2.ContainerError(code="code", message="message")
pb_str = _base64.b64encode(pb.SerializeToString())
obj = proto_type.from_string(pb_str)
assert obj.to_python_std().code == "code"
assert obj.to_python_std().message == "message"
def test_proto():
proto_type = _proto.create_protobuf(_errors_pb2.ContainerError)
assert proto_type.short_class_string(
) == "Types.Proto(flyteidl.core.errors_pb2.ContainerError)"
pb = _errors_pb2.ContainerError(code="code", message="message")
obj = proto_type.from_python_std(pb)
obj2 = proto_type.from_flyte_idl(obj.to_flyte_idl())
assert obj == obj2
obj = obj.to_python_std()
obj2 = obj2.to_python_std()
assert obj.code == "code"
assert obj.message == "message"
assert obj2.code == "code"
assert obj2.message == "message"
def _proto_sdk_type_from_tag(tag):
"""
:param Text tag:
:rtype: _proto.Protobuf
"""
if "." not in tag:
raise _user_exceptions.FlyteValueException(
tag, "Protobuf tag must include at least one '.' to delineate package and object name.",
)
module, name = tag.rsplit(".", 1)
try:
pb_module = _importer.import_module(module)
except ImportError:
raise _user_exceptions.FlyteAssertion(
"Could not resolve the protobuf definition @ {}. Is the protobuf library installed?".format(module)
)
if not hasattr(pb_module, name):
raise _user_exceptions.FlyteAssertion("Could not find the protobuf named: {} @ {}.".format(name, module))
return _proto.create_protobuf(getattr(pb_module, name))
def _proto_sdk_type_from_tag(tag):
"""
:param Text tag:
:rtype: _proto.Protobuf
"""
return _proto.create_protobuf(_load_type_from_tag(tag))
def test_wrong_type():
with _pytest.raises(_user_exceptions.FlyteTypeException):
_proto.create_protobuf(int)
