text: # -*- coding: utf-8 -*-
  
  # Copyright 2022 Google LLC
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #     http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  #
  
  import uuid
  
  from google.cloud import aiplatform
  
  from tests.system.aiplatform import e2e_base
  
  # project
  _TEST_INDEX_DISPLAY_NAME = "index_display_name"
  _TEST_INDEX_DESCRIPTION = "index_description"
  _TEST_INDEX_DISTANCE_MEASURE_TYPE = "SQUARED_L2_DISTANCE"
  
  _TEST_INDEX_CONFIG_DIMENSIONS = 100
  _TEST_INDEX_APPROXIMATE_NEIGHBORS_COUNT = 150
  _TEST_LEAF_NODE_EMBEDDING_COUNT = 123
  _TEST_LEAF_NODES_TO_SEARCH_PERCENT = 50
  
  
  _TEST_CONTENTS_DELTA_URI = (
      "gs://cloud-samples-data-us-central1/vertex-ai/matching_engine/glove100/initial"
  )
  _TEST_CONTENTS_DELTA_URI_UPDATE = (
      "gs://cloud-samples-data-us-central1/vertex-ai/matching_engine/glove100/incremental"
  )
  _TEST_IS_COMPLETE_OVERWRITE = True
  _TEST_INDEX_DISTANCE_MEASURE_TYPE = "SQUARED_L2_DISTANCE"
  
  
  _TEST_LABELS = {"my_key": "my_value"}
  _TEST_DISPLAY_NAME_UPDATE = "my new display name"
  _TEST_DESCRIPTION_UPDATE = "my description update"
  _TEST_LABELS_UPDATE = {"my_key_update": "my_value_update"}
  
  # ENDPOINT
  _TEST_INDEX_ENDPOINT_DISPLAY_NAME = "endpoint_name"
  _TEST_INDEX_ENDPOINT_DESCRIPTION = "my endpoint"
  
  _TEST_INDEX_ENDPOINT_VPC_NETWORK = "projects/{}/global/networks/{}".format(
      e2e_base._PROJECT_NUMBER, e2e_base._VPC_NETWORK_NAME
  )
  
  # DEPLOYED INDEX
  _TEST_DEPLOYED_INDEX_ID = f"deployed_index_id_{uuid.uuid4()}"
  _TEST_DEPLOYED_INDEX_DISPLAY_NAME = f"deployed_index_display_name_{uuid.uuid4()}"
  _TEST_MIN_REPLICA_COUNT_UPDATED = 4
  _TEST_MAX_REPLICA_COUNT_UPDATED = 4
  
  # QUERY
  _TEST_MATCH_QUERY = query = [
      -0.11333,
      0.48402,
      0.090771,
      -0.22439,
      0.034206,
      -0.55831,
      0.041849,
      -0.53573,
      0.18809,
      -0.58722,
      0.015313,
      -0.014555,
      0.80842,
      -0.038519,
      0.75348,
      0.70502,
      -0.17863,
      0.3222,
      0.67575,
      0.67198,
      0.26044,
      0.4187,
      -0.34122,
      0.2286,
      -0.53529,
      1.2582,
      -0.091543,
      0.19716,
      -0.037454,
      -0.3336,
      0.31399,
      0.36488,
      0.71263,
      0.1307,
      -0.24654,
      -0.52445,
      -0.036091,
      0.55068,
      0.10017,
      0.48095,
      0.71104,
      -0.053462,
      0.22325,
      0.30917,
      -0.39926,
      0.036634,
      -0.35431,
      -0.42795,
      0.46444,
      0.25586,
      0.68257,
      -0.20821,
      0.38433,
      0.055773,
      -0.2539,
      -0.20804,
      0.52522,
      -0.11399,
      -0.3253,
      -0.44104,
      0.17528,
      0.62255,
      0.50237,
      -0.7607,
      -0.071786,
      0.0080131,
      -0.13286,
      0.50097,
      0.18824,
      -0.54722,
      -0.42664,
      0.4292,
      0.14877,
      -0.0072514,
      -0.16484,
      -0.059798,
      0.9895,
      -0.61738,
      0.054169,
      0.48424,
      -0.35084,
      -0.27053,
      0.37829,
      0.11503,
      -0.39613,
      0.24266,
      0.39147,
      -0.075256,
      0.65093,
      -0.20822,
      -0.17456,
      0.53571,
      -0.16537,
      0.13582,
      -0.56016,
      0.016964,
      0.1277,
      0.94071,
      -0.22608,
      -0.021106,
  ]
  
  
  class TestMatchingEngine(e2e_base.TestEndToEnd):
  
      _temp_prefix = "temp_vertex_sdk_e2e_matching_engine_test"
  
      def test_create_get_list_matching_engine_index(self, shared_state):
          aiplatform.init(
              project=e2e_base._PROJECT,
              location=e2e_base._LOCATION,
          )
  
          # Create an index
          index = aiplatform.MatchingEngineIndex.create_tree_ah_index(
              display_name=_TEST_INDEX_DISPLAY_NAME,
              contents_delta_uri=_TEST_CONTENTS_DELTA_URI,
              dimensions=_TEST_INDEX_CONFIG_DIMENSIONS,
              approximate_neighbors_count=_TEST_INDEX_APPROXIMATE_NEIGHBORS_COUNT,
              distance_measure_type=_TEST_INDEX_DISTANCE_MEASURE_TYPE,
              leaf_node_embedding_count=_TEST_LEAF_NODE_EMBEDDING_COUNT,
              leaf_nodes_to_search_percent=_TEST_LEAF_NODES_TO_SEARCH_PERCENT,
              description=_TEST_INDEX_DESCRIPTION,
              labels=_TEST_LABELS,
          )
  
          shared_state["resources"] = [index]
          shared_state["index"] = index
          shared_state["index_name"] = index.resource_name
  
          # Verify that the retrieved index is the same
          get_index = aiplatform.MatchingEngineIndex(index_name=index.resource_name)
          assert index.resource_name == get_index.resource_name
  
          # Create index and check that it is listed
          list_indexes = aiplatform.MatchingEngineIndex.list()
          assert get_index.resource_name in [
              index.resource_name for index in list_indexes
          ]
  
          # Update the index metadata
          updated_index = get_index.update_metadata(
              display_name=_TEST_DISPLAY_NAME_UPDATE,
              description=_TEST_DESCRIPTION_UPDATE,
              labels=_TEST_LABELS_UPDATE,
          )
  
          assert updated_index.name == get_index.name
          # TODO: Reinstate assertions once b/220005272 is fixed.
          # assert updated_index.display_name == _TEST_DISPLAY_NAME_UPDATE
          # assert updated_index.description == _TEST_DESCRIPTION_UPDATE
          # assert updated_index.labels == _TEST_LABELS_UPDATE
  
          # Update the index embeddings
          updated_index = get_index.update_embeddings(
              contents_delta_uri=_TEST_CONTENTS_DELTA_URI_UPDATE,
              is_complete_overwrite=_TEST_IS_COMPLETE_OVERWRITE,
          )
  
          assert updated_index.name == get_index.name
  
          # Delete index and check that it is no longer listed
          index.delete()
          list_indexes = aiplatform.MatchingEngineIndex.list()
          assert get_index.resource_name not in [
              index.resource_name for index in list_indexes
          ]
  