def testConvertNamedSignatureToSignatureDef(self):
signatures_proto = manifest_pb2.Signatures()
generic_signature = manifest_pb2.GenericSignature()
generic_signature.map["input_key"].CopyFrom(
manifest_pb2.TensorBinding(tensor_name="input"))
signatures_proto.named_signatures[
signature_constants.PREDICT_INPUTS].generic_signature.CopyFrom(
generic_signature)
generic_signature = manifest_pb2.GenericSignature()
generic_signature.map["output_key"].CopyFrom(
manifest_pb2.TensorBinding(tensor_name="output"))
signatures_proto.named_signatures[
signature_constants.PREDICT_OUTPUTS].generic_signature.CopyFrom(
generic_signature)
signature_def = bundle_shim._convert_named_signatures_to_signature_def(
signatures_proto)
self.assertEqual(signature_def.method_name,
signature_constants.PREDICT_METHOD_NAME)
self.assertEqual(len(signature_def.inputs), 1)
self.assertEqual(len(signature_def.outputs), 1)
self.assertProtoEquals(signature_def.inputs["input_key"],
meta_graph_pb2.TensorInfo(name="input"))
self.assertProtoEquals(signature_def.outputs["output_key"],
meta_graph_pb2.TensorInfo(name="output"))
def testConvertDefaultSignatureBadTypeToSignatureDef(self):
signatures_proto = manifest_pb2.Signatures()
generic_signature = manifest_pb2.GenericSignature()
signatures_proto.default_signature.generic_signature.CopyFrom(
generic_signature)
with self.assertRaises(RuntimeError) as cm:
_ = bundle_shim._convert_default_signature_to_signature_def(
signatures_proto)
def testConvertDefaultSignatureGenericToSignatureDef(self):
signatures_proto = manifest_pb2.Signatures()
generic_signature = manifest_pb2.GenericSignature()
signatures_proto.default_signature.generic_signature.CopyFrom(
generic_signature)
signature_def = bundle_shim._convert_default_signature_to_signature_def(
signatures_proto)
self.assertEquals(signature_def, None)