@@ -1061,6 +1061,130 @@ def test_create_agent_engine(
10611061 retry = _TEST_RETRY ,
10621062 )
10631063
1064+ def test_create_agent_engine_with_protobuf_agent_card (
1065+ self ,
1066+ create_agent_engine_mock ,
1067+ cloud_storage_create_bucket_mock ,
1068+ tarfile_open_mock ,
1069+ cloudpickle_dump_mock ,
1070+ cloudpickle_load_mock ,
1071+ importlib_metadata_version_mock ,
1072+ get_agent_engine_mock ,
1073+ get_gca_resource_mock ,
1074+ ):
1075+ class CapitalizeEngineWithCard (CapitalizeEngine ):
1076+ def __init__ (self , card ):
1077+ self .agent_card = card
1078+
1079+ from google .protobuf import struct_pb2
1080+
1081+ card = struct_pb2 .Struct ()
1082+ card ["name" ] = "test_agent_card"
1083+ agent = CapitalizeEngineWithCard (card )
1084+
1085+ agent_engines .create (
1086+ agent ,
1087+ display_name = _TEST_AGENT_ENGINE_DISPLAY_NAME ,
1088+ requirements = _TEST_AGENT_ENGINE_REQUIREMENTS ,
1089+ extra_packages = [_TEST_AGENT_ENGINE_EXTRA_PACKAGE_PATH ],
1090+ )
1091+
1092+ expected_reasoning_engine = types .ReasoningEngine (
1093+ display_name = _TEST_AGENT_ENGINE_DISPLAY_NAME ,
1094+ spec = types .ReasoningEngineSpec (
1095+ package_spec = _TEST_AGENT_ENGINE_PACKAGE_SPEC ,
1096+ agent_framework = _agent_engines ._DEFAULT_AGENT_FRAMEWORK ,
1097+ ),
1098+ )
1099+ from google .protobuf import json_format
1100+
1101+ expected_class_method = struct_pb2 .Struct ()
1102+ expected_class_method .CopyFrom (_TEST_AGENT_ENGINE_QUERY_SCHEMA )
1103+ expected_class_method ["a2a_agent_card" ] = json_format .MessageToJson (card )
1104+ expected_reasoning_engine .spec .class_methods .append (expected_class_method )
1105+
1106+ create_agent_engine_mock .assert_called_with (
1107+ parent = _TEST_PARENT ,
1108+ reasoning_engine = expected_reasoning_engine ,
1109+ )
1110+
1111+ def test_create_agent_engine_with_pydantic_agent_card (
1112+ self ,
1113+ create_agent_engine_mock ,
1114+ cloud_storage_create_bucket_mock ,
1115+ tarfile_open_mock ,
1116+ cloudpickle_dump_mock ,
1117+ cloudpickle_load_mock ,
1118+ importlib_metadata_version_mock ,
1119+ get_agent_engine_mock ,
1120+ get_gca_resource_mock ,
1121+ ):
1122+ class CapitalizeEngineWithCard (CapitalizeEngine ):
1123+ def __init__ (self , card ):
1124+ self .agent_card = card
1125+
1126+ import pydantic
1127+
1128+ class DummyPydanticCard (pydantic .BaseModel ):
1129+ name : str = "test_pydantic_card"
1130+
1131+ card = DummyPydanticCard ()
1132+ agent = CapitalizeEngineWithCard (card )
1133+
1134+ agent_engines .create (
1135+ agent ,
1136+ display_name = _TEST_AGENT_ENGINE_DISPLAY_NAME ,
1137+ requirements = _TEST_AGENT_ENGINE_REQUIREMENTS ,
1138+ extra_packages = [_TEST_AGENT_ENGINE_EXTRA_PACKAGE_PATH ],
1139+ )
1140+
1141+ expected_reasoning_engine = types .ReasoningEngine (
1142+ display_name = _TEST_AGENT_ENGINE_DISPLAY_NAME ,
1143+ spec = types .ReasoningEngineSpec (
1144+ package_spec = _TEST_AGENT_ENGINE_PACKAGE_SPEC ,
1145+ agent_framework = _agent_engines ._DEFAULT_AGENT_FRAMEWORK ,
1146+ ),
1147+ )
1148+ from google .protobuf import struct_pb2
1149+
1150+ expected_class_method = struct_pb2 .Struct ()
1151+ expected_class_method .CopyFrom (_TEST_AGENT_ENGINE_QUERY_SCHEMA )
1152+ expected_class_method ["a2a_agent_card" ] = card .model_dump_json ()
1153+ expected_reasoning_engine .spec .class_methods .append (expected_class_method )
1154+
1155+ create_agent_engine_mock .assert_called_with (
1156+ parent = _TEST_PARENT ,
1157+ reasoning_engine = expected_reasoning_engine ,
1158+ )
1159+
1160+ def test_create_agent_engine_with_invalid_agent_card (
1161+ self ,
1162+ create_agent_engine_mock ,
1163+ cloud_storage_create_bucket_mock ,
1164+ tarfile_open_mock ,
1165+ cloudpickle_dump_mock ,
1166+ cloudpickle_load_mock ,
1167+ importlib_metadata_version_mock ,
1168+ get_agent_engine_mock ,
1169+ get_gca_resource_mock ,
1170+ ):
1171+ class CapitalizeEngineWithCard (CapitalizeEngine ):
1172+ def __init__ (self , card ):
1173+ self .agent_card = card
1174+
1175+ agent = CapitalizeEngineWithCard (card = "invalid_card_type_string" )
1176+
1177+ with pytest .raises (
1178+ ValueError ,
1179+ match = "Failed to serialize agent_card: Unsupported AgentCard type" ,
1180+ ):
1181+ agent_engines .create (
1182+ agent ,
1183+ display_name = _TEST_AGENT_ENGINE_DISPLAY_NAME ,
1184+ requirements = _TEST_AGENT_ENGINE_REQUIREMENTS ,
1185+ extra_packages = [_TEST_AGENT_ENGINE_EXTRA_PACKAGE_PATH ],
1186+ )
1187+
10641188 def test_create_agent_engine_requirements_from_file (
10651189 self ,
10661190 create_agent_engine_mock ,
0 commit comments