def test_update_index_endpoint(self, update_index_endpoint_mock):
aiplatform.init(project=_TEST_PROJECT)
my_index_endpoint = aiplatform.MatchingEngineIndexEndpoint(
index_endpoint_name=_TEST_INDEX_ENDPOINT_ID)
updated_endpoint = my_index_endpoint.update(
display_name=_TEST_DISPLAY_NAME_UPDATE,
description=_TEST_DESCRIPTION_UPDATE,
labels=_TEST_LABELS_UPDATE,
request_metadata=_TEST_REQUEST_METADATA,
)
expected = gca_index_endpoint.IndexEndpoint(
name=_TEST_INDEX_ENDPOINT_NAME,
display_name=_TEST_DISPLAY_NAME_UPDATE,
description=_TEST_DESCRIPTION_UPDATE,
labels=_TEST_LABELS_UPDATE,
)
update_index_endpoint_mock.assert_called_once_with(
index_endpoint=expected,
update_mask=field_mask_pb2.FieldMask(
paths=["labels", "display_name", "description"]),
metadata=_TEST_REQUEST_METADATA,
)
assert updated_endpoint.gca_resource == expected
def test_init_index_endpoint(self, index_endpoint_name,
get_index_endpoint_mock):
aiplatform.init(project=_TEST_PROJECT)
my_index_endpoint = aiplatform.MatchingEngineIndexEndpoint(
index_endpoint_name=index_endpoint_name)
get_index_endpoint_mock.assert_called_once_with(
name=my_index_endpoint.resource_name, retry=base._DEFAULT_RETRY)
def test_deploy_index(self, deploy_index_mock, undeploy_index_mock):
aiplatform.init(project=_TEST_PROJECT)
my_index_endpoint = aiplatform.MatchingEngineIndexEndpoint(
index_endpoint_name=_TEST_INDEX_ENDPOINT_ID)
# Get index
my_index = aiplatform.MatchingEngineIndex(index_name=_TEST_INDEX_NAME)
my_index_endpoint = my_index_endpoint.deploy_index(
index=my_index,
deployed_index_id=_TEST_DEPLOYED_INDEX_ID,
display_name=_TEST_DEPLOYED_INDEX_DISPLAY_NAME,
min_replica_count=_TEST_MIN_REPLICA_COUNT,
max_replica_count=_TEST_MAX_REPLICA_COUNT,
enable_access_logging=_TEST_ENABLE_ACCESS_LOGGING,
reserved_ip_ranges=_TEST_RESERVED_IP_RANGES,
deployment_group=_TEST_DEPLOYMENT_GROUP,
auth_config_audiences=_TEST_AUTH_CONFIG_AUDIENCES,
auth_config_allowed_issuers=_TEST_AUTH_CONFIG_ALLOWED_ISSUERS,
request_metadata=_TEST_REQUEST_METADATA,
)
deploy_index_mock.assert_called_once_with(
index_endpoint=my_index_endpoint.resource_name,
deployed_index=gca_matching_engine_index_endpoint.DeployedIndex(
id=_TEST_DEPLOYED_INDEX_ID,
index=my_index.resource_name,
display_name=_TEST_DEPLOYED_INDEX_DISPLAY_NAME,
enable_access_logging=_TEST_ENABLE_ACCESS_LOGGING,
reserved_ip_ranges=_TEST_RESERVED_IP_RANGES,
deployment_group=_TEST_DEPLOYMENT_GROUP,
automatic_resources={
"min_replica_count": _TEST_MIN_REPLICA_COUNT,
"max_replica_count": _TEST_MAX_REPLICA_COUNT,
},
deployed_index_auth_config=gca_matching_engine_index_endpoint.
DeployedIndexAuthConfig(
auth_provider=gca_matching_engine_index_endpoint.
DeployedIndexAuthConfig.AuthProvider(
audiences=_TEST_AUTH_CONFIG_AUDIENCES,
allowed_issuers=_TEST_AUTH_CONFIG_ALLOWED_ISSUERS,
)),
),
metadata=_TEST_REQUEST_METADATA,
)
my_index_endpoint = my_index_endpoint.undeploy_index(
deployed_index_id=_TEST_DEPLOYED_INDEX_ID)
undeploy_index_mock.assert_called_once_with(
index_endpoint=my_index_endpoint.resource_name,
deployed_index_id=_TEST_DEPLOYED_INDEX_ID,
metadata=_TEST_REQUEST_METADATA,
)
def test_delete_index_endpoint(self, delete_index_endpoint_mock, sync):
aiplatform.init(project=_TEST_PROJECT)
my_index_endpoint = aiplatform.MatchingEngineIndexEndpoint(
index_endpoint_name=_TEST_INDEX_ENDPOINT_ID)
my_index_endpoint.delete(sync=sync)
if not sync:
my_index_endpoint.wait()
delete_index_endpoint_mock.assert_called_once_with(
name=my_index_endpoint.resource_name)
def test_delete_index_endpoint_with_force(self, undeploy_index_mock,
delete_index_endpoint_mock,
sync):
my_index_endpoint = aiplatform.MatchingEngineIndexEndpoint(
index_endpoint_name=_TEST_INDEX_ENDPOINT_NAME)
my_index_endpoint.delete(force=True, sync=sync)
if not sync:
my_index_endpoint.wait()
# undeploy_index_mock should be called if force is set to True
assert undeploy_index_mock.call_count == 2
delete_index_endpoint_mock.assert_called_once_with(
name=_TEST_INDEX_ENDPOINT_NAME)
def test_delete_index_endpoint_without_force(self, undeploy_index_mock,
delete_index_endpoint_mock,
sync):
my_index_endpoint = aiplatform.MatchingEngineIndexEndpoint(
index_endpoint_name=_TEST_INDEX_ENDPOINT_NAME)
my_index_endpoint.delete(sync=sync)
if not sync:
my_index_endpoint.wait()
# undeploy_index_mock should not be called unless force is set to True
undeploy_index_mock.assert_not_called()
delete_index_endpoint_mock.assert_called_once_with(
name=_TEST_INDEX_ENDPOINT_NAME)
def test_mutate_deployed_index(self, mutate_deployed_index_mock):
aiplatform.init(project=_TEST_PROJECT)
my_index_endpoint = aiplatform.MatchingEngineIndexEndpoint(
index_endpoint_name=_TEST_INDEX_ENDPOINT_ID)
my_index_endpoint.mutate_deployed_index(
deployed_index_id=_TEST_DEPLOYED_INDEX_ID,
min_replica_count=_TEST_MIN_REPLICA_COUNT_UPDATED,
max_replica_count=_TEST_MAX_REPLICA_COUNT_UPDATED,
request_metadata=_TEST_REQUEST_METADATA,
)
mutate_deployed_index_mock.assert_called_once_with(
index_endpoint=_TEST_INDEX_ENDPOINT_NAME,
deployed_index=gca_matching_engine_index_endpoint.DeployedIndex(
id=_TEST_DEPLOYED_INDEX_ID,
automatic_resources={
"min_replica_count": _TEST_MIN_REPLICA_COUNT_UPDATED,
"max_replica_count": _TEST_MAX_REPLICA_COUNT_UPDATED,
},
),
metadata=_TEST_REQUEST_METADATA,
)