Skip to content

Commit 1b5ae44

Browse files
vertex-sdk-botcopybara-github
authored andcommitted
feat: add support for return public endpoint dns name in matching engine
PiperOrigin-RevId: 525507137
1 parent 4b0722c commit 1b5ae44

File tree

2 files changed

+70
-1
lines changed

2 files changed

+70
-1
lines changed

google/cloud/aiplatform/matching_engine/matching_engine_index_endpoint.py

+6
Original file line numberDiff line numberDiff line change
@@ -344,6 +344,12 @@ def _create(
344344

345345
return index_obj
346346

347+
@property
348+
def public_endpoint_domain_name(self) -> Optional[str]:
349+
"""Public endpoint DNS name."""
350+
self._assert_gca_resource_is_available()
351+
return self._gca_resource.public_endpoint_domain_name
352+
347353
def update(
348354
self,
349355
display_name: str,

tests/unit/aiplatform/test_matching_engine_index_endpoint.py

+64-1
Original file line numberDiff line numberDiff line change
@@ -60,6 +60,9 @@
6060

6161
# index_endpoint
6262
_TEST_INDEX_ENDPOINT_ID = "index_endpoint_id"
63+
_TEST_INDEX_ENDPOINT_PUBLIC_DNS = (
64+
"1114627793.us-central1-249381615684.vdb.vertexai.goog"
65+
)
6366
_TEST_INDEX_ENDPOINT_NAME = f"{_TEST_PARENT}/indexEndpoints/{_TEST_INDEX_ENDPOINT_ID}"
6467
_TEST_INDEX_ENDPOINT_DISPLAY_NAME = "index_endpoint_display_name"
6568
_TEST_INDEX_ENDPOINT_DESCRIPTION = "index_endpoint_description"
@@ -308,6 +311,57 @@ def get_index_endpoint_mock():
308311
yield get_index_endpoint_mock
309312

310313

314+
@pytest.fixture
315+
def get_index_public_endpoint_mock():
316+
with patch.object(
317+
index_endpoint_service_client.IndexEndpointServiceClient, "get_index_endpoint"
318+
) as get_index_public_endpoint_mock:
319+
index_endpoint = gca_index_endpoint.IndexEndpoint(
320+
name=_TEST_INDEX_ENDPOINT_NAME,
321+
display_name=_TEST_INDEX_ENDPOINT_DISPLAY_NAME,
322+
description=_TEST_INDEX_ENDPOINT_DESCRIPTION,
323+
public_endpoint_domain_name=_TEST_INDEX_ENDPOINT_PUBLIC_DNS,
324+
)
325+
index_endpoint.deployed_indexes = [
326+
gca_index_endpoint.DeployedIndex(
327+
id=_TEST_DEPLOYED_INDEX_ID,
328+
index=_TEST_INDEX_NAME,
329+
display_name=_TEST_DEPLOYED_INDEX_DISPLAY_NAME,
330+
enable_access_logging=_TEST_ENABLE_ACCESS_LOGGING,
331+
deployment_group=_TEST_DEPLOYMENT_GROUP,
332+
automatic_resources={
333+
"min_replica_count": _TEST_MIN_REPLICA_COUNT,
334+
"max_replica_count": _TEST_MAX_REPLICA_COUNT,
335+
},
336+
deployed_index_auth_config=gca_index_endpoint.DeployedIndexAuthConfig(
337+
auth_provider=gca_index_endpoint.DeployedIndexAuthConfig.AuthProvider(
338+
audiences=_TEST_AUTH_CONFIG_AUDIENCES,
339+
allowed_issuers=_TEST_AUTH_CONFIG_ALLOWED_ISSUERS,
340+
)
341+
),
342+
),
343+
gca_index_endpoint.DeployedIndex(
344+
id=f"{_TEST_DEPLOYED_INDEX_ID}_2",
345+
index=f"{_TEST_INDEX_NAME}_2",
346+
display_name=_TEST_DEPLOYED_INDEX_DISPLAY_NAME,
347+
enable_access_logging=_TEST_ENABLE_ACCESS_LOGGING,
348+
deployment_group=_TEST_DEPLOYMENT_GROUP,
349+
automatic_resources={
350+
"min_replica_count": _TEST_MIN_REPLICA_COUNT,
351+
"max_replica_count": _TEST_MAX_REPLICA_COUNT,
352+
},
353+
deployed_index_auth_config=gca_index_endpoint.DeployedIndexAuthConfig(
354+
auth_provider=gca_index_endpoint.DeployedIndexAuthConfig.AuthProvider(
355+
audiences=_TEST_AUTH_CONFIG_AUDIENCES,
356+
allowed_issuers=_TEST_AUTH_CONFIG_ALLOWED_ISSUERS,
357+
)
358+
),
359+
),
360+
]
361+
get_index_public_endpoint_mock.return_value = index_endpoint
362+
yield get_index_public_endpoint_mock
363+
364+
311365
@pytest.fixture
312366
def deploy_index_mock():
313367
with patch.object(
@@ -556,7 +610,7 @@ def test_create_index_endpoint_with_network_init(self, create_index_endpoint_moc
556610
metadata=_TEST_REQUEST_METADATA,
557611
)
558612

559-
@pytest.mark.usefixtures("get_index_endpoint_mock")
613+
@pytest.mark.usefixtures("get_index_public_endpoint_mock")
560614
def test_create_index_endpoint_with_public_endpoint_enabled(
561615
self, create_index_endpoint_mock
562616
):
@@ -569,6 +623,10 @@ def test_create_index_endpoint_with_public_endpoint_enabled(
569623
labels=_TEST_LABELS,
570624
)
571625

626+
my_index_endpoint = aiplatform.MatchingEngineIndexEndpoint(
627+
index_endpoint_name=_TEST_INDEX_ENDPOINT_ID
628+
)
629+
572630
expected = gca_index_endpoint.IndexEndpoint(
573631
display_name=_TEST_INDEX_ENDPOINT_DISPLAY_NAME,
574632
description=_TEST_INDEX_ENDPOINT_DESCRIPTION,
@@ -582,6 +640,11 @@ def test_create_index_endpoint_with_public_endpoint_enabled(
582640
metadata=_TEST_REQUEST_METADATA,
583641
)
584642

643+
assert (
644+
my_index_endpoint.public_endpoint_domain_name
645+
== _TEST_INDEX_ENDPOINT_PUBLIC_DNS
646+
)
647+
585648
def test_create_index_endpoint_missing_argument_throw_error(
586649
self, create_index_endpoint_mock
587650
):

0 commit comments

Comments
 (0)