Skip to content

Commit 75ba7c1

Browse files
shawn-yang-googlecopybara-github
shawn-yang-google
authored and
copybara-github
committed
feat: Support Developer Connect in AE
PiperOrigin-RevId: 832480435
1 parent 2c362b3 commit 75ba7c1

File tree

5 files changed

+410
-9
lines changed

5 files changed

+410
-9
lines changed

‎tests/unit/vertexai/genai/test_agent_engines.py‎

Lines changed: 222 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1062,6 +1062,47 @@ def test_create_agent_engine_config_with_source_packages(
10621062
== _TEST_AGENT_ENGINE_IDENTITY_TYPE_SERVICE_ACCOUNT
10631063
)
10641064

1065+
def test_create_agent_engine_config_with_developer_connect_source(self):
1066+
with tempfile.TemporaryDirectory() as tmpdir:
1067+
requirements_file_path = os.path.join(tmpdir, "requirements.txt")
1068+
with open(requirements_file_path, "w") as f:
1069+
f.write("requests==2.0.0")
1070+
developer_connect_source = {
1071+
"git_repository_link": "projects/test-project/locations/us-central1/connections/test-connection/gitRepositoryLinks/test-repo",
1072+
"revision": "main",
1073+
"dir": "agent",
1074+
}
1075+
config = self.client.agent_engines._create_config(
1076+
mode="create",
1077+
display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME,
1078+
description=_TEST_AGENT_ENGINE_DESCRIPTION,
1079+
developer_connect_source=developer_connect_source,
1080+
entrypoint_module="main",
1081+
entrypoint_object="app",
1082+
requirements_file=requirements_file_path,
1083+
class_methods=_TEST_AGENT_ENGINE_CLASS_METHODS,
1084+
agent_framework=_TEST_AGENT_FRAMEWORK,
1085+
identity_type=_TEST_AGENT_ENGINE_IDENTITY_TYPE_SERVICE_ACCOUNT,
1086+
python_version=_TEST_PYTHON_VERSION_OVERRIDE,
1087+
)
1088+
assert config["display_name"] == _TEST_AGENT_ENGINE_DISPLAY_NAME
1089+
assert config["description"] == _TEST_AGENT_ENGINE_DESCRIPTION
1090+
assert config["spec"]["agent_framework"] == _TEST_AGENT_FRAMEWORK
1091+
assert config["spec"]["source_code_spec"] == {
1092+
"developer_connect_source": developer_connect_source,
1093+
"python_spec": {
1094+
"version": _TEST_PYTHON_VERSION_OVERRIDE,
1095+
"entrypoint_module": "main",
1096+
"entrypoint_object": "app",
1097+
"requirements_file": requirements_file_path,
1098+
},
1099+
}
1100+
assert config["spec"]["class_methods"] == _TEST_AGENT_ENGINE_CLASS_METHODS
1101+
assert (
1102+
config["spec"]["identity_type"]
1103+
== _TEST_AGENT_ENGINE_IDENTITY_TYPE_SERVICE_ACCOUNT
1104+
)
1105+
10651106
@mock.patch.object(
10661107
_agent_engines_utils,
10671108
"_create_base64_encoded_tarball",
@@ -1521,6 +1562,89 @@ def test_prepare_without_creds(
15211562
gcs_dir_name=_TEST_GCS_DIR_NAME,
15221563
)
15231564

1565+
@pytest.mark.parametrize(
1566+
"operation_name,resource_name,expected_id,expected_exception,expected_message",
1567+
[
1568+
(
1569+
f"projects/123/locations/us-central1/reasoningEngines/{_TEST_RESOURCE_ID}/operations/456",
1570+
"",
1571+
_TEST_RESOURCE_ID,
1572+
None,
1573+
None,
1574+
),
1575+
(
1576+
"",
1577+
f"projects/123/locations/us-central1/reasoningEngines/{_TEST_RESOURCE_ID}",
1578+
_TEST_RESOURCE_ID,
1579+
None,
1580+
None,
1581+
),
1582+
(
1583+
"projects/123/locations/us-central1/reasoningEngines/other_id/operations/456",
1584+
f"projects/123/locations/us-central1/reasoningEngines/{_TEST_RESOURCE_ID}",
1585+
_TEST_RESOURCE_ID,
1586+
None,
1587+
None,
1588+
),
1589+
(
1590+
"",
1591+
"",
1592+
None,
1593+
ValueError,
1594+
"Resource name or operation name cannot be empty.",
1595+
),
1596+
(
1597+
"invalid/operation/name",
1598+
"",
1599+
None,
1600+
ValueError,
1601+
"Failed to parse reasoning engine ID from operation name",
1602+
),
1603+
(
1604+
f"projects/123/locations/us-central1/reasoningEngines/{_TEST_RESOURCE_ID}",
1605+
"",
1606+
None,
1607+
ValueError,
1608+
"Failed to parse reasoning engine ID from operation name",
1609+
),
1610+
(
1611+
"",
1612+
"invalid/resource/name",
1613+
None,
1614+
ValueError,
1615+
"Failed to parse reasoning engine ID from resource name",
1616+
),
1617+
(
1618+
"",
1619+
f"projects/123/locations/us-central1/reasoningEngines/{_TEST_RESOURCE_ID}/operations/456",
1620+
None,
1621+
ValueError,
1622+
"Failed to parse reasoning engine ID from resource name",
1623+
),
1624+
],
1625+
)
1626+
def test_get_reasoning_engine_id(
1627+
self,
1628+
operation_name,
1629+
resource_name,
1630+
expected_id,
1631+
expected_exception,
1632+
expected_message,
1633+
):
1634+
if expected_exception:
1635+
with pytest.raises(expected_exception) as excinfo:
1636+
_agent_engines_utils._get_reasoning_engine_id(
1637+
operation_name=operation_name, resource_name=resource_name
1638+
)
1639+
assert expected_message in str(excinfo.value)
1640+
else:
1641+
assert (
1642+
_agent_engines_utils._get_reasoning_engine_id(
1643+
operation_name=operation_name, resource_name=resource_name
1644+
)
1645+
== expected_id
1646+
)
1647+
15241648

15251649
@pytest.mark.usefixtures("google_auth_mock")
15261650
class TestAgentEngine:
@@ -1570,7 +1694,18 @@ def test_list_agent_engine(self):
15701694
@pytest.mark.usefixtures("caplog")
15711695
@mock.patch.object(_agent_engines_utils, "_prepare")
15721696
@mock.patch.object(_agent_engines_utils, "_await_operation")
1573-
def test_create_agent_engine(self, mock_await_operation, mock_prepare, caplog):
1697+
@mock.patch.object(
1698+
_agent_engines_utils,
1699+
"_get_reasoning_engine_id",
1700+
return_value=_TEST_RESOURCE_ID,
1701+
)
1702+
def test_create_agent_engine(
1703+
self,
1704+
mock_get_reasoning_engine_id,
1705+
mock_await_operation,
1706+
mock_prepare,
1707+
caplog,
1708+
):
15741709
mock_await_operation.return_value = _genai_types.AgentEngineOperation(
15751710
response=_genai_types.ReasoningEngine(
15761711
name=_TEST_AGENT_ENGINE_RESOURCE_NAME,
@@ -1620,8 +1755,14 @@ def test_create_agent_engine(self, mock_await_operation, mock_prepare, caplog):
16201755

16211756
@mock.patch.object(agent_engines.AgentEngines, "_create_config")
16221757
@mock.patch.object(_agent_engines_utils, "_await_operation")
1758+
@mock.patch.object(
1759+
_agent_engines_utils,
1760+
"_get_reasoning_engine_id",
1761+
return_value=_TEST_RESOURCE_ID,
1762+
)
16231763
def test_create_agent_engine_lightweight(
16241764
self,
1765+
mock_get_reasoning_engine_id,
16251766
mock_await_operation,
16261767
mock_create_config,
16271768
):
@@ -1657,8 +1798,14 @@ def test_create_agent_engine_lightweight(
16571798

16581799
@mock.patch.object(agent_engines.AgentEngines, "_create_config")
16591800
@mock.patch.object(_agent_engines_utils, "_await_operation")
1801+
@mock.patch.object(
1802+
_agent_engines_utils,
1803+
"_get_reasoning_engine_id",
1804+
return_value=_TEST_RESOURCE_ID,
1805+
)
16601806
def test_create_agent_engine_with_env_vars_dict(
16611807
self,
1808+
mock_get_reasoning_engine_id,
16621809
mock_await_operation,
16631810
mock_create_config,
16641811
):
@@ -1718,6 +1865,7 @@ def test_create_agent_engine_with_env_vars_dict(
17181865
labels=None,
17191866
class_methods=None,
17201867
source_packages=None,
1868+
developer_connect_source=None,
17211869
entrypoint_module=None,
17221870
entrypoint_object=None,
17231871
requirements_file=None,
@@ -1748,8 +1896,14 @@ def test_create_agent_engine_with_env_vars_dict(
17481896

17491897
@mock.patch.object(agent_engines.AgentEngines, "_create_config")
17501898
@mock.patch.object(_agent_engines_utils, "_await_operation")
1899+
@mock.patch.object(
1900+
_agent_engines_utils,
1901+
"_get_reasoning_engine_id",
1902+
return_value=_TEST_RESOURCE_ID,
1903+
)
17511904
def test_create_agent_engine_with_custom_service_account(
17521905
self,
1906+
mock_get_reasoning_engine_id,
17531907
mock_await_operation,
17541908
mock_create_config,
17551909
):
@@ -1812,6 +1966,7 @@ def test_create_agent_engine_with_custom_service_account(
18121966
agent_server_mode=None,
18131967
class_methods=None,
18141968
source_packages=None,
1969+
developer_connect_source=None,
18151970
entrypoint_module=None,
18161971
entrypoint_object=None,
18171972
requirements_file=None,
@@ -1842,8 +1997,14 @@ def test_create_agent_engine_with_custom_service_account(
18421997

18431998
@mock.patch.object(agent_engines.AgentEngines, "_create_config")
18441999
@mock.patch.object(_agent_engines_utils, "_await_operation")
2000+
@mock.patch.object(
2001+
_agent_engines_utils,
2002+
"_get_reasoning_engine_id",
2003+
return_value=_TEST_RESOURCE_ID,
2004+
)
18452005
def test_create_agent_engine_with_experimental_mode(
18462006
self,
2007+
mock_get_reasoning_engine_id,
18472008
mock_await_operation,
18482009
mock_create_config,
18492010
):
@@ -1905,6 +2066,7 @@ def test_create_agent_engine_with_experimental_mode(
19052066
agent_server_mode=_genai_types.AgentServerMode.EXPERIMENTAL,
19062067
class_methods=None,
19072068
source_packages=None,
2069+
developer_connect_source=None,
19082070
entrypoint_module=None,
19092071
entrypoint_object=None,
19102072
requirements_file=None,
@@ -1939,8 +2101,14 @@ def test_create_agent_engine_with_experimental_mode(
19392101
return_value="test_tarball",
19402102
)
19412103
@mock.patch.object(_agent_engines_utils, "_await_operation")
2104+
@mock.patch.object(
2105+
_agent_engines_utils,
2106+
"_get_reasoning_engine_id",
2107+
return_value=_TEST_RESOURCE_ID,
2108+
)
19422109
def test_create_agent_engine_with_source_packages(
19432110
self,
2111+
mock_get_reasoning_engine_id,
19442112
mock_await_operation,
19452113
mock_create_base64_encoded_tarball,
19462114
):
@@ -2001,8 +2169,14 @@ def test_create_agent_engine_with_source_packages(
20012169

20022170
@mock.patch.object(agent_engines.AgentEngines, "_create_config")
20032171
@mock.patch.object(_agent_engines_utils, "_await_operation")
2172+
@mock.patch.object(
2173+
_agent_engines_utils,
2174+
"_get_reasoning_engine_id",
2175+
return_value=_TEST_RESOURCE_ID,
2176+
)
20042177
def test_create_agent_engine_with_class_methods(
20052178
self,
2179+
mock_get_reasoning_engine_id,
20062180
mock_await_operation,
20072181
mock_create_config,
20082182
):
@@ -2061,6 +2235,7 @@ def test_create_agent_engine_with_class_methods(
20612235
agent_server_mode=None,
20622236
class_methods=_TEST_AGENT_ENGINE_CLASS_METHODS,
20632237
source_packages=None,
2238+
developer_connect_source=None,
20642239
entrypoint_module=None,
20652240
entrypoint_object=None,
20662241
requirements_file=None,
@@ -2088,8 +2263,14 @@ def test_create_agent_engine_with_class_methods(
20882263

20892264
@mock.patch.object(agent_engines.AgentEngines, "_create_config")
20902265
@mock.patch.object(_agent_engines_utils, "_await_operation")
2266+
@mock.patch.object(
2267+
_agent_engines_utils,
2268+
"_get_reasoning_engine_id",
2269+
return_value=_TEST_RESOURCE_ID,
2270+
)
20912271
def test_create_agent_engine_with_agent_framework(
20922272
self,
2273+
mock_get_reasoning_engine_id,
20932274
mock_await_operation,
20942275
mock_create_config,
20952276
):
@@ -2148,6 +2329,7 @@ def test_create_agent_engine_with_agent_framework(
21482329
agent_server_mode=None,
21492330
class_methods=None,
21502331
source_packages=None,
2332+
developer_connect_source=None,
21512333
entrypoint_module=None,
21522334
entrypoint_object=None,
21532335
requirements_file=None,
@@ -2696,8 +2878,17 @@ def test_operation_schemas(
26962878
@mock.patch.object(_agent_engines_utils, "_prepare")
26972879
@mock.patch.object(agent_engines.AgentEngines, "_create")
26982880
@mock.patch.object(_agent_engines_utils, "_await_operation")
2881+
@mock.patch.object(
2882+
_agent_engines_utils,
2883+
"_get_reasoning_engine_id",
2884+
return_value=_TEST_RESOURCE_ID,
2885+
)
26992886
def test_create_agent_engine_with_creds(
2700-
self, mock_await_operation, mock_create, mock_prepare
2887+
self,
2888+
mock_get_reasoning_engine_id,
2889+
mock_await_operation,
2890+
mock_create,
2891+
mock_prepare,
27012892
):
27022893
mock_operation = mock.Mock()
27032894
mock_operation.name = _TEST_AGENT_ENGINE_OPERATION_NAME
@@ -2728,8 +2919,18 @@ def test_create_agent_engine_with_creds(
27282919
@mock.patch.object(agent_engines.AgentEngines, "_create")
27292920
@mock.patch("google.auth.default")
27302921
@mock.patch.object(_agent_engines_utils, "_await_operation")
2922+
@mock.patch.object(
2923+
_agent_engines_utils,
2924+
"_get_reasoning_engine_id",
2925+
return_value=_TEST_RESOURCE_ID,
2926+
)
27312927
def test_create_agent_engine_without_creds(
2732-
self, mock_await_operation, mock_auth_default, mock_create, mock_prepare
2928+
self,
2929+
mock_get_reasoning_engine_id,
2930+
mock_await_operation,
2931+
mock_auth_default,
2932+
mock_create,
2933+
mock_prepare,
27332934
):
27342935
mock_operation = mock.Mock()
27352936
mock_operation.name = _TEST_AGENT_ENGINE_OPERATION_NAME
@@ -2765,8 +2966,17 @@ def test_create_agent_engine_without_creds(
27652966
@mock.patch.object(_agent_engines_utils, "_prepare")
27662967
@mock.patch.object(agent_engines.AgentEngines, "_create")
27672968
@mock.patch.object(_agent_engines_utils, "_await_operation")
2969+
@mock.patch.object(
2970+
_agent_engines_utils,
2971+
"_get_reasoning_engine_id",
2972+
return_value=_TEST_RESOURCE_ID,
2973+
)
27682974
def test_create_agent_engine_with_no_creds_in_client(
2769-
self, mock_await_operation, mock_create, mock_prepare
2975+
self,
2976+
mock_get_reasoning_engine_id,
2977+
mock_await_operation,
2978+
mock_create,
2979+
mock_prepare,
27702980
):
27712981
mock_operation = mock.Mock()
27722982
mock_operation.name = _TEST_AGENT_ENGINE_OPERATION_NAME
@@ -2812,7 +3022,14 @@ def setup_method(self):
28123022

28133023
@mock.patch.object(_agent_engines_utils, "_prepare")
28143024
@mock.patch.object(_agent_engines_utils, "_await_operation")
2815-
def test_create_agent_engine_error(self, mock_await_operation, mock_prepare):
3025+
@mock.patch.object(
3026+
_agent_engines_utils,
3027+
"_get_reasoning_engine_id",
3028+
return_value=_TEST_RESOURCE_ID,
3029+
)
3030+
def test_create_agent_engine_error(
3031+
self, mock_get_reasoning_engine_id, mock_await_operation, mock_prepare
3032+
):
28163033
mock_await_operation.return_value = _genai_types.AgentEngineOperation(
28173034
error=_TEST_AGENT_ENGINE_ERROR,
28183035
)

0 commit comments

Comments
 (0)