Skip to content

Commit 65bf9b6

Browse files
yeesiancopybara-github
authored andcommitted
fix: Rename agent_engine parameter to agent in Agent Engines.
PiperOrigin-RevId: 796691005
1 parent 26477ce commit 65bf9b6

File tree

4 files changed

+192
-161
lines changed

4 files changed

+192
-161
lines changed

tests/unit/vertexai/genai/test_agent_engines.py

Lines changed: 27 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -790,7 +790,7 @@ def test_create_agent_engine_config_lightweight(self, mock_prepare):
790790
def test_create_agent_engine_config_full(self, mock_prepare):
791791
config = self.client.agent_engines._create_config(
792792
mode="create",
793-
agent_engine=self.test_agent,
793+
agent=self.test_agent,
794794
staging_bucket=_TEST_STAGING_BUCKET,
795795
requirements=_TEST_AGENT_ENGINE_REQUIREMENTS,
796796
display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME,
@@ -839,7 +839,7 @@ def test_create_agent_engine_config_full(self, mock_prepare):
839839
def test_update_agent_engine_config_full(self, mock_prepare):
840840
config = self.client.agent_engines._create_config(
841841
mode="update",
842-
agent_engine=self.test_agent,
842+
agent=self.test_agent,
843843
staging_bucket=_TEST_STAGING_BUCKET,
844844
requirements=_TEST_AGENT_ENGINE_REQUIREMENTS,
845845
display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME,
@@ -939,7 +939,7 @@ def test_await_operation(self):
939939

940940
def test_register_api_methods(self):
941941
agent = self.client.agent_engines._register_api_methods(
942-
agent=_genai_types.AgentEngine(
942+
agent_engine=_genai_types.AgentEngine(
943943
api_client=self.client.agent_engines._api_client,
944944
api_resource=_genai_types.ReasoningEngine(
945945
spec=_genai_types.ReasoningEngineSpec(
@@ -956,7 +956,9 @@ def test_register_api_methods(self):
956956

957957
@pytest.mark.usefixtures("caplog")
958958
def test_invalid_requirement_warning(self, caplog):
959-
_agent_engines_utils._parse_constraints(["invalid requirement line"])
959+
_agent_engines_utils._parse_constraints(
960+
constraints=["invalid requirement line"],
961+
)
960962
assert "Failed to parse constraint" in caplog.text
961963

962964
def test_requirements_with_whl_files(self):
@@ -965,7 +967,9 @@ def test_requirements_with_whl_files(self):
965967
"/content/wxPython-4.2.3-cp39-cp39-macosx_12_0_x86_64.whl",
966968
"https://wxpython.org/Phoenix/snapshot-builds/wxPython-4.2.2-cp38-cp38-macosx_12_0_x86_64.whl",
967969
]
968-
result = _agent_engines_utils._parse_constraints(whl_files)
970+
result = _agent_engines_utils._parse_constraints(
971+
constraints=whl_files,
972+
)
969973
assert result == {
970974
"wxPython-4.2.2-cp38-cp38-macosx_12_0_x86_64.whl": None,
971975
"wxPython-4.2.3-cp39-cp39-macosx_12_0_x86_64.whl": None,
@@ -975,7 +979,10 @@ def test_requirements_with_whl_files(self):
975979
def test_compare_requirements_with_required_packages(self):
976980
requirements = {"requests": "2.0.0"}
977981
constraints = ["requests==1.0.0"]
978-
result = _agent_engines_utils._compare_requirements(requirements, constraints)
982+
result = _agent_engines_utils._compare_requirements(
983+
requirements=requirements,
984+
constraints=constraints,
985+
)
979986
assert result == {
980987
"actions": {"append": set()},
981988
"warnings": {
@@ -989,7 +996,7 @@ def test_scan_simple_object(self):
989996
"""Test scanning an object importing a known third-party package."""
990997
fake_obj = _create_fake_object_with_module("requests")
991998
requirements = _agent_engines_utils._scan_requirements(
992-
fake_obj,
999+
obj=fake_obj,
9931000
package_distributions=_TEST_PACKAGE_DISTRIBUTIONS,
9941001
)
9951002
assert requirements == {
@@ -1003,7 +1010,7 @@ def test_scan_object_with_stdlib_module(self):
10031010
"""Test that stdlib modules are ignored by default."""
10041011
fake_obj_stdlib = _create_fake_object_with_module("json")
10051012
requirements = _agent_engines_utils._scan_requirements(
1006-
fake_obj_stdlib,
1013+
obj=fake_obj_stdlib,
10071014
package_distributions=_TEST_PACKAGE_DISTRIBUTIONS,
10081015
)
10091016
# Requirements should not contain 'json',
@@ -1024,7 +1031,7 @@ def test_scan_with_default_ignore_modules(self, monkeypatch):
10241031
set(original_base) | {"requests"},
10251032
)
10261033
requirements = _agent_engines_utils._scan_requirements(
1027-
fake_obj,
1034+
obj=fake_obj,
10281035
package_distributions=_TEST_PACKAGE_DISTRIBUTIONS,
10291036
)
10301037
# Requirements should not contain 'requests',
@@ -1039,7 +1046,7 @@ def test_scan_with_explicit_ignore_modules(self):
10391046
"""Test explicitly ignoring a module."""
10401047
fake_obj = _create_fake_object_with_module("requests")
10411048
requirements = _agent_engines_utils._scan_requirements(
1042-
fake_obj,
1049+
obj=fake_obj,
10431050
ignore_modules=["requests"],
10441051
package_distributions=_TEST_PACKAGE_DISTRIBUTIONS,
10451052
)
@@ -1152,7 +1159,7 @@ def test_create_agent_engine(self, mock_await_operation, mock_prepare, caplog):
11521159
) as request_mock:
11531160
request_mock.return_value = genai_types.HttpResponse(body="")
11541161
self.client.agent_engines.create(
1155-
agent_engine=self.test_agent,
1162+
agent=self.test_agent,
11561163
config=_genai_types.AgentEngineConfig(
11571164
display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME,
11581165
description=_TEST_AGENT_ENGINE_DESCRIPTION,
@@ -1263,7 +1270,7 @@ def test_create_agent_engine_with_env_vars_dict(
12631270
) as request_mock:
12641271
request_mock.return_value = genai_types.HttpResponse(body="")
12651272
self.client.agent_engines.create(
1266-
agent_engine=self.test_agent,
1273+
agent=self.test_agent,
12671274
config=_genai_types.AgentEngineConfig(
12681275
display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME,
12691276
requirements=_TEST_AGENT_ENGINE_REQUIREMENTS,
@@ -1274,7 +1281,7 @@ def test_create_agent_engine_with_env_vars_dict(
12741281
)
12751282
mock_create_config.assert_called_with(
12761283
mode="create",
1277-
agent_engine=self.test_agent,
1284+
agent=self.test_agent,
12781285
staging_bucket=_TEST_STAGING_BUCKET,
12791286
requirements=_TEST_AGENT_ENGINE_REQUIREMENTS,
12801287
display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME,
@@ -1327,7 +1334,7 @@ def test_update_agent_engine_requirements(
13271334
request_mock.return_value = genai_types.HttpResponse(body="")
13281335
self.client.agent_engines.update(
13291336
name=_TEST_AGENT_ENGINE_RESOURCE_NAME,
1330-
agent_engine=self.test_agent,
1337+
agent=self.test_agent,
13311338
config=_genai_types.AgentEngineConfig(
13321339
staging_bucket=_TEST_STAGING_BUCKET,
13331340
requirements=_TEST_AGENT_ENGINE_REQUIREMENTS,
@@ -1383,7 +1390,7 @@ def test_update_agent_engine_extra_packages(
13831390
request_mock.return_value = genai_types.HttpResponse(body="")
13841391
self.client.agent_engines.update(
13851392
name=_TEST_AGENT_ENGINE_RESOURCE_NAME,
1386-
agent_engine=self.test_agent,
1393+
agent=self.test_agent,
13871394
config=_genai_types.AgentEngineConfig(
13881395
staging_bucket=_TEST_STAGING_BUCKET,
13891396
requirements=_TEST_AGENT_ENGINE_REQUIREMENTS,
@@ -1437,7 +1444,7 @@ def test_update_agent_engine_env_vars(
14371444
request_mock.return_value = genai_types.HttpResponse(body="")
14381445
self.client.agent_engines.update(
14391446
name=_TEST_AGENT_ENGINE_RESOURCE_NAME,
1440-
agent_engine=self.test_agent,
1447+
agent=self.test_agent,
14411448
config=_genai_types.AgentEngineConfig(
14421449
staging_bucket=_TEST_STAGING_BUCKET,
14431450
requirements=_TEST_AGENT_ENGINE_REQUIREMENTS,
@@ -1570,7 +1577,7 @@ def test_query_agent_engine(self):
15701577
) as request_mock:
15711578
request_mock.return_value = genai_types.HttpResponse(body="")
15721579
agent = self.client.agent_engines._register_api_methods(
1573-
agent=_genai_types.AgentEngine(
1580+
agent_engine=_genai_types.AgentEngine(
15741581
api_client=self.client.agent_engines,
15751582
api_resource=_genai_types.ReasoningEngine(
15761583
name=_TEST_AGENT_ENGINE_RESOURCE_NAME,
@@ -1596,7 +1603,7 @@ def test_query_agent_engine(self):
15961603

15971604
def test_query_agent_engine_async(self):
15981605
agent = self.client.agent_engines._register_api_methods(
1599-
agent=_genai_types.AgentEngine(
1606+
agent_engine=_genai_types.AgentEngine(
16001607
api_async_client=agent_engines.AsyncAgentEngines(
16011608
api_client_=self.client.agent_engines._api_client
16021609
),
@@ -1631,7 +1638,7 @@ def test_query_agent_engine_stream(self):
16311638
self.client.agent_engines._api_client, "request_streamed"
16321639
) as request_mock:
16331640
agent = self.client.agent_engines._register_api_methods(
1634-
agent=_genai_types.AgentEngine(
1641+
agent_engine=_genai_types.AgentEngine(
16351642
api_client=self.client.agent_engines,
16361643
api_resource=_genai_types.ReasoningEngine(
16371644
name=_TEST_AGENT_ENGINE_RESOURCE_NAME,
@@ -1664,7 +1671,7 @@ async def mock_async_generator():
16641671
) as request_mock:
16651672
request_mock.return_value = mock_async_generator()
16661673
agent = self.client.agent_engines._register_api_methods(
1667-
agent=_genai_types.AgentEngine(
1674+
agent_engine=_genai_types.AgentEngine(
16681675
api_client=self.client.agent_engines,
16691676
api_resource=_genai_types.ReasoningEngine(
16701677
name=_TEST_AGENT_ENGINE_RESOURCE_NAME,

0 commit comments

Comments
 (0)