Skip to content

Commit 2dccde8

Browse files
vertex-sdk-botcopybara-github
vertex-sdk-bot
authored and
copybara-github
committed
feat: Allow setting ingress_control_config when creating a SandboxEnvironmentTemplate via the SDK (Python, Java, JS).
PiperOrigin-RevId: 974166963
1 parent 2c60077 commit 2dccde8

4 files changed

Lines changed: 202 additions & 104 deletions

File tree

‎agentplatform/_genai/sandbox_templates.py‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -62,6 +62,13 @@ def _CreateSandboxEnvironmentTemplateConfig_to_vertex(
6262
getv(from_object, ["egress_control_config"]),
6363
)
6464

65+
if getv(from_object, ["ingress_control_config"]) is not None:
66+
setv(
67+
parent_object,
68+
["ingressControlConfig"],
69+
getv(from_object, ["ingress_control_config"]),
70+
)
71+
6572
return to_object
6673

6774

‎agentplatform/_genai/types/__init__.py‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -3457,15 +3457,15 @@
34573457
"SandboxEnvironmentTemplateEgressControlConfig",
34583458
"SandboxEnvironmentTemplateEgressControlConfigDict",
34593459
"SandboxEnvironmentTemplateEgressControlConfigOrDict",
3460-
"CreateSandboxEnvironmentTemplateConfig",
3461-
"CreateSandboxEnvironmentTemplateConfigDict",
3462-
"CreateSandboxEnvironmentTemplateConfigOrDict",
34633460
"PSCAutomationConfig",
34643461
"PSCAutomationConfigDict",
34653462
"PSCAutomationConfigOrDict",
34663463
"PrivateServiceConnectConfig",
34673464
"PrivateServiceConnectConfigDict",
34683465
"PrivateServiceConnectConfigOrDict",
3466+
"CreateSandboxEnvironmentTemplateConfig",
3467+
"CreateSandboxEnvironmentTemplateConfigDict",
3468+
"CreateSandboxEnvironmentTemplateConfigOrDict",
34693469
"SandboxEnvironmentTemplate",
34703470
"SandboxEnvironmentTemplateDict",
34713471
"SandboxEnvironmentTemplateOrDict",

‎agentplatform/_genai/types/common.py‎

Lines changed: 99 additions & 92 deletions
Original file line numberDiff line numberDiff line change
@@ -17603,98 +17603,6 @@ class SandboxEnvironmentTemplateEgressControlConfigDict(TypedDict, total=False):
1760317603
]
1760417604

1760517605

17606-
class CreateSandboxEnvironmentTemplateConfig(_common.BaseModel):
17607-
"""Config for creating a Sandbox Template."""
17608-
17609-
http_options: Optional[genai_types.HttpOptions] = Field(
17610-
default=None, description="""Used to override HTTP request options."""
17611-
)
17612-
wait_for_completion: Optional[bool] = Field(
17613-
default=True,
17614-
description="""Waits for the operation to complete before returning.""",
17615-
)
17616-
custom_container_environment: Optional[
17617-
SandboxEnvironmentTemplateCustomContainerEnvironment
17618-
] = Field(
17619-
default=None,
17620-
description="""The custom container environment for the sandbox template.""",
17621-
)
17622-
default_container_environment: Optional[
17623-
SandboxEnvironmentTemplateDefaultContainerEnvironment
17624-
] = Field(
17625-
default=None,
17626-
description="""The default container environment for the sandbox template.""",
17627-
)
17628-
egress_control_config: Optional[SandboxEnvironmentTemplateEgressControlConfig] = (
17629-
Field(
17630-
default=None,
17631-
description="""The egress control config for the sandbox template.""",
17632-
)
17633-
)
17634-
17635-
17636-
class CreateSandboxEnvironmentTemplateConfigDict(TypedDict, total=False):
17637-
"""Config for creating a Sandbox Template."""
17638-
17639-
http_options: Optional[genai_types.HttpOptions]
17640-
"""Used to override HTTP request options."""
17641-
17642-
wait_for_completion: Optional[bool]
17643-
"""Waits for the operation to complete before returning."""
17644-
17645-
custom_container_environment: Optional[
17646-
SandboxEnvironmentTemplateCustomContainerEnvironmentDict
17647-
]
17648-
"""The custom container environment for the sandbox template."""
17649-
17650-
default_container_environment: Optional[
17651-
SandboxEnvironmentTemplateDefaultContainerEnvironmentDict
17652-
]
17653-
"""The default container environment for the sandbox template."""
17654-
17655-
egress_control_config: Optional[SandboxEnvironmentTemplateEgressControlConfigDict]
17656-
"""The egress control config for the sandbox template."""
17657-
17658-
17659-
CreateSandboxEnvironmentTemplateConfigOrDict = Union[
17660-
CreateSandboxEnvironmentTemplateConfig, CreateSandboxEnvironmentTemplateConfigDict
17661-
]
17662-
17663-
17664-
class _CreateSandboxEnvironmentTemplateRequestParameters(_common.BaseModel):
17665-
"""Parameters for creating Sandbox Environment Templates."""
17666-
17667-
name: Optional[str] = Field(
17668-
default=None,
17669-
description="""Name of the agent runtime to create the template under.""",
17670-
)
17671-
display_name: Optional[str] = Field(
17672-
default=None, description="""The display name of the sandbox template."""
17673-
)
17674-
config: Optional[CreateSandboxEnvironmentTemplateConfig] = Field(
17675-
default=None, description=""""""
17676-
)
17677-
17678-
17679-
class _CreateSandboxEnvironmentTemplateRequestParametersDict(TypedDict, total=False):
17680-
"""Parameters for creating Sandbox Environment Templates."""
17681-
17682-
name: Optional[str]
17683-
"""Name of the agent runtime to create the template under."""
17684-
17685-
display_name: Optional[str]
17686-
"""The display name of the sandbox template."""
17687-
17688-
config: Optional[CreateSandboxEnvironmentTemplateConfigDict]
17689-
""""""
17690-
17691-
17692-
_CreateSandboxEnvironmentTemplateRequestParametersOrDict = Union[
17693-
_CreateSandboxEnvironmentTemplateRequestParameters,
17694-
_CreateSandboxEnvironmentTemplateRequestParametersDict,
17695-
]
17696-
17697-
1769817606
class PSCAutomationConfig(_common.BaseModel):
1769917607
"""PSC config that is used to automatically create PSC endpoints in the user projects."""
1770017608

@@ -17798,6 +17706,105 @@ class PrivateServiceConnectConfigDict(TypedDict, total=False):
1779817706
]
1779917707

1780017708

17709+
class CreateSandboxEnvironmentTemplateConfig(_common.BaseModel):
17710+
"""Config for creating a Sandbox Template."""
17711+
17712+
http_options: Optional[genai_types.HttpOptions] = Field(
17713+
default=None, description="""Used to override HTTP request options."""
17714+
)
17715+
wait_for_completion: Optional[bool] = Field(
17716+
default=True,
17717+
description="""Waits for the operation to complete before returning.""",
17718+
)
17719+
custom_container_environment: Optional[
17720+
SandboxEnvironmentTemplateCustomContainerEnvironment
17721+
] = Field(
17722+
default=None,
17723+
description="""The custom container environment for the sandbox template.""",
17724+
)
17725+
default_container_environment: Optional[
17726+
SandboxEnvironmentTemplateDefaultContainerEnvironment
17727+
] = Field(
17728+
default=None,
17729+
description="""The default container environment for the sandbox template.""",
17730+
)
17731+
egress_control_config: Optional[SandboxEnvironmentTemplateEgressControlConfig] = (
17732+
Field(
17733+
default=None,
17734+
description="""The egress control config for the sandbox template.""",
17735+
)
17736+
)
17737+
ingress_control_config: Optional[PrivateServiceConnectConfig] = Field(
17738+
default=None,
17739+
description="""The ingress control config for the sandbox template.""",
17740+
)
17741+
17742+
17743+
class CreateSandboxEnvironmentTemplateConfigDict(TypedDict, total=False):
17744+
"""Config for creating a Sandbox Template."""
17745+
17746+
http_options: Optional[genai_types.HttpOptions]
17747+
"""Used to override HTTP request options."""
17748+
17749+
wait_for_completion: Optional[bool]
17750+
"""Waits for the operation to complete before returning."""
17751+
17752+
custom_container_environment: Optional[
17753+
SandboxEnvironmentTemplateCustomContainerEnvironmentDict
17754+
]
17755+
"""The custom container environment for the sandbox template."""
17756+
17757+
default_container_environment: Optional[
17758+
SandboxEnvironmentTemplateDefaultContainerEnvironmentDict
17759+
]
17760+
"""The default container environment for the sandbox template."""
17761+
17762+
egress_control_config: Optional[SandboxEnvironmentTemplateEgressControlConfigDict]
17763+
"""The egress control config for the sandbox template."""
17764+
17765+
ingress_control_config: Optional[PrivateServiceConnectConfigDict]
17766+
"""The ingress control config for the sandbox template."""
17767+
17768+
17769+
CreateSandboxEnvironmentTemplateConfigOrDict = Union[
17770+
CreateSandboxEnvironmentTemplateConfig, CreateSandboxEnvironmentTemplateConfigDict
17771+
]
17772+
17773+
17774+
class _CreateSandboxEnvironmentTemplateRequestParameters(_common.BaseModel):
17775+
"""Parameters for creating Sandbox Environment Templates."""
17776+
17777+
name: Optional[str] = Field(
17778+
default=None,
17779+
description="""Name of the agent runtime to create the template under.""",
17780+
)
17781+
display_name: Optional[str] = Field(
17782+
default=None, description="""The display name of the sandbox template."""
17783+
)
17784+
config: Optional[CreateSandboxEnvironmentTemplateConfig] = Field(
17785+
default=None, description=""""""
17786+
)
17787+
17788+
17789+
class _CreateSandboxEnvironmentTemplateRequestParametersDict(TypedDict, total=False):
17790+
"""Parameters for creating Sandbox Environment Templates."""
17791+
17792+
name: Optional[str]
17793+
"""Name of the agent runtime to create the template under."""
17794+
17795+
display_name: Optional[str]
17796+
"""The display name of the sandbox template."""
17797+
17798+
config: Optional[CreateSandboxEnvironmentTemplateConfigDict]
17799+
""""""
17800+
17801+
17802+
_CreateSandboxEnvironmentTemplateRequestParametersOrDict = Union[
17803+
_CreateSandboxEnvironmentTemplateRequestParameters,
17804+
_CreateSandboxEnvironmentTemplateRequestParametersDict,
17805+
]
17806+
17807+
1780117808
class SandboxEnvironmentTemplate(_common.BaseModel):
1780217809
"""A sandbox environment template."""
1780317810

‎tests/unit/agentplatform/genai/test_sandbox.py‎

Lines changed: 93 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -130,12 +130,10 @@ def test_generate_browser_ws_headers(
130130
mock_http_client.request.return_value = genai_types.HttpResponse(
131131
body=b'{"endpoint": "test/endpoint"}', headers={}
132132
)
133-
ws_url, headers = (
134-
self.client.sandboxes.generate_browser_ws_headers(
135-
sandbox_environment=mock_sandbox,
136-
service_account_email=_TEST_SERVICE_ACCOUNT_EMAIL,
137-
timeout=3600,
138-
)
133+
ws_url, headers = self.client.sandboxes.generate_browser_ws_headers(
134+
sandbox_environment=mock_sandbox,
135+
service_account_email=_TEST_SERVICE_ACCOUNT_EMAIL,
136+
timeout=3600,
139137
)
140138
assert ws_url == "wss://test-us-central1.example.vertexai.goog/test/endpoint"
141139
assert (
@@ -297,6 +295,94 @@ def test_create_without_spec_template_or_snapshot_raises(self, mock_create):
297295
mock_create.assert_not_called()
298296

299297

298+
@pytest.mark.usefixtures("google_auth_mock")
299+
class TestSandboxTemplates:
300+
"""Tests for sandbox_templates.SandboxTemplates.
301+
302+
Exercises the create-side wiring of fields on
303+
`CreateSandboxEnvironmentTemplateConfig`, including
304+
`ingress_control_config` (see cl/972355621), which is otherwise not
305+
covered by the sandboxes.create tests above.
306+
"""
307+
308+
def setup_method(self):
309+
importlib.reload(initializer)
310+
importlib.reload(aiplatform)
311+
importlib.reload(agentplatform)
312+
self.client = agentplatform.Client(
313+
project=_TEST_PROJECT,
314+
location=_TEST_LOCATION,
315+
credentials=_TEST_CREDENTIALS,
316+
)
317+
318+
def teardown_method(self):
319+
initializer.global_pool.shutdown(wait=True)
320+
321+
def test_create_config_ingress_control_config_is_optional(self):
322+
"""Constructing without ingress_control_config must succeed."""
323+
config = agentplatform_types.CreateSandboxEnvironmentTemplateConfig()
324+
assert config.ingress_control_config is None
325+
326+
def test_create_config_accepts_ingress_control_config_object(self):
327+
"""The typed field accepts a PrivateServiceConnectConfig instance."""
328+
psc = agentplatform_types.PrivateServiceConnectConfig(
329+
enable_private_service_connect=True,
330+
project_allowlist=["test-project"],
331+
)
332+
config = agentplatform_types.CreateSandboxEnvironmentTemplateConfig(
333+
ingress_control_config=psc,
334+
)
335+
assert config.ingress_control_config is not None
336+
assert config.ingress_control_config.enable_private_service_connect is True
337+
assert config.ingress_control_config.project_allowlist == ["test-project"]
338+
339+
def test_create_config_accepts_ingress_control_config_dict(self):
340+
"""The typed field validates and coerces a dict input."""
341+
config = (
342+
agentplatform_types.CreateSandboxEnvironmentTemplateConfig.model_validate(
343+
{
344+
"ingress_control_config": {
345+
"enable_private_service_connect": True,
346+
"project_allowlist": ["test-project"],
347+
},
348+
}
349+
)
350+
)
351+
assert isinstance(
352+
config.ingress_control_config,
353+
agentplatform_types.PrivateServiceConnectConfig,
354+
)
355+
assert config.ingress_control_config.enable_private_service_connect is True
356+
assert config.ingress_control_config.project_allowlist == ["test-project"]
357+
358+
@mock.patch.object(sandbox_templates.SandboxTemplates, "_create")
359+
def test_create_forwards_ingress_control_config(self, mock_create):
360+
"""templates.create(...) passes ingress_control_config through to _create."""
361+
mock_create.return_value = mock.Mock()
362+
363+
config = agentplatform_types.CreateSandboxEnvironmentTemplateConfig(
364+
ingress_control_config=agentplatform_types.PrivateServiceConnectConfig(
365+
enable_private_service_connect=True,
366+
project_allowlist=["test-project"],
367+
),
368+
wait_for_completion=False,
369+
)
370+
self.client.sandboxes.templates.create(
371+
name=_TEST_AGENT_ENGINE_RESOURCE_NAME,
372+
display_name="test-template",
373+
config=config,
374+
)
375+
376+
mock_create.assert_called_once()
377+
_, kwargs = mock_create.call_args
378+
assert kwargs["name"] == _TEST_AGENT_ENGINE_RESOURCE_NAME
379+
assert kwargs["display_name"] == "test-template"
380+
ingress = kwargs["config"].ingress_control_config
381+
assert ingress is not None
382+
assert ingress.enable_private_service_connect is True
383+
assert ingress.project_allowlist == ["test-project"]
384+
385+
300386
_MODULES = pytest.mark.parametrize(
301387
"module",
302388
[sandboxes, vertexai_sandboxes],
@@ -412,9 +498,7 @@ def test_generate_access_token_retries_transient_failures(module, status_code):
412498
[transient, transient, _ok_response()],
413499
)
414500

415-
with default_patch, session_patch, mock.patch.object(
416-
module.time, "sleep"
417-
) as sleep:
501+
with default_patch, session_patch, mock.patch.object(module.time, "sleep") as sleep:
418502
client_obj = module.Sandboxes(api_client_=mock.Mock())
419503
token = client_obj.generate_access_token(
420504
service_account_email=_TEST_SERVICE_ACCOUNT_EMAIL

0 commit comments

Comments
 (0)